1"""Interface to the compiler's internal symbol tables"""
2
3import _symtable
4from _symtable import (USE, DEF_GLOBAL, DEF_LOCAL, DEF_PARAM,
5     DEF_IMPORT, DEF_BOUND, OPT_IMPORT_STAR, OPT_EXEC, OPT_BARE_EXEC,
6     SCOPE_OFF, SCOPE_MASK, FREE, GLOBAL_IMPLICIT, GLOBAL_EXPLICIT, CELL, LOCAL)
7
8import weakref
9
10__all__ = ["symtable", "SymbolTable", "Class", "Function", "Symbol"]
11
12def symtable(code, filename, compile_type):
13    raw = _symtable.symtable(code, filename, compile_type)
14    for top in raw.itervalues():
15        if top.name == 'top':
16            break
17    return _newSymbolTable(top, filename)
18
19class SymbolTableFactory:
20    def __init__(self):
21        self.__memo = weakref.WeakValueDictionary()
22
23    def new(self, table, filename):
24        if table.type == _symtable.TYPE_FUNCTION:
25            return Function(table, filename)
26        if table.type == _symtable.TYPE_CLASS:
27            return Class(table, filename)
28        return SymbolTable(table, filename)
29
30    def __call__(self, table, filename):
31        key = table, filename
32        obj = self.__memo.get(key, None)
33        if obj is None:
34            obj = self.__memo[key] = self.new(table, filename)
35        return obj
36
37_newSymbolTable = SymbolTableFactory()
38
39
40class SymbolTable(object):
41
42    def __init__(self, raw_table, filename):
43        self._table = raw_table
44        self._filename = filename
45        self._symbols = {}
46
47    def __repr__(self):
48        if self.__class__ == SymbolTable:
49            kind = ""
50        else:
51            kind = "%s " % self.__class__.__name__
52
53        if self._table.name == "global":
54            return "<{0}SymbolTable for module {1}>".format(kind, self._filename)
55        else:
56            return "<{0}SymbolTable for {1} in {2}>".format(kind,
57                                                            self._table.name,
58                                                            self._filename)
59
60    def get_type(self):
61        if self._table.type == _symtable.TYPE_MODULE:
62            return "module"
63        if self._table.type == _symtable.TYPE_FUNCTION:
64            return "function"
65        if self._table.type == _symtable.TYPE_CLASS:
66            return "class"
67        assert self._table.type in (1, 2, 3), \
68               "unexpected type: {0}".format(self._table.type)
69
70    def get_id(self):
71        return self._table.id
72
73    def get_name(self):
74        return self._table.name
75
76    def get_lineno(self):
77        return self._table.lineno
78
79    def is_optimized(self):
80        return bool(self._table.type == _symtable.TYPE_FUNCTION
81                    and not self._table.optimized)
82
83    def is_nested(self):
84        return bool(self._table.nested)
85
86    def has_children(self):
87        return bool(self._table.children)
88
89    def has_exec(self):
90        """Return true if the scope uses exec"""
91        return bool(self._table.optimized & (OPT_EXEC | OPT_BARE_EXEC))
92
93    def has_import_star(self):
94        """Return true if the scope uses import *"""
95        return bool(self._table.optimized & OPT_IMPORT_STAR)
96
97    def get_identifiers(self):
98        return self._table.symbols.keys()
99
100    def lookup(self, name):
101        sym = self._symbols.get(name)
102        if sym is None:
103            flags = self._table.symbols[name]
104            namespaces = self.__check_children(name)
105            sym = self._symbols[name] = Symbol(name, flags, namespaces)
106        return sym
107
108    def get_symbols(self):
109        return [self.lookup(ident) for ident in self.get_identifiers()]
110
111    def __check_children(self, name):
112        return [_newSymbolTable(st, self._filename)
113                for st in self._table.children
114                if st.name == name]
115
116    def get_children(self):
117        return [_newSymbolTable(st, self._filename)
118                for st in self._table.children]
119
120
121class Function(SymbolTable):
122
123    # Default values for instance variables
124    __params = None
125    __locals = None
126    __frees = None
127    __globals = None
128
129    def __idents_matching(self, test_func):
130        return tuple([ident for ident in self.get_identifiers()
131                      if test_func(self._table.symbols[ident])])
132
133    def get_parameters(self):
134        if self.__params is None:
135            self.__params = self.__idents_matching(lambda x:x & DEF_PARAM)
136        return self.__params
137
138    def get_locals(self):
139        if self.__locals is None:
140            locs = (LOCAL, CELL)
141            test = lambda x: ((x >> SCOPE_OFF) & SCOPE_MASK) in locs
142            self.__locals = self.__idents_matching(test)
143        return self.__locals
144
145    def get_globals(self):
146        if self.__globals is None:
147            glob = (GLOBAL_IMPLICIT, GLOBAL_EXPLICIT)
148            test = lambda x:((x >> SCOPE_OFF) & SCOPE_MASK) in glob
149            self.__globals = self.__idents_matching(test)
150        return self.__globals
151
152    def get_frees(self):
153        if self.__frees is None:
154            is_free = lambda x:((x >> SCOPE_OFF) & SCOPE_MASK) == FREE
155            self.__frees = self.__idents_matching(is_free)
156        return self.__frees
157
158
159class Class(SymbolTable):
160
161    __methods = None
162
163    def get_methods(self):
164        if self.__methods is None:
165            d = {}
166            for st in self._table.children:
167                d[st.name] = 1
168            self.__methods = tuple(d)
169        return self.__methods
170
171
172class Symbol(object):
173
174    def __init__(self, name, flags, namespaces=None):
175        self.__name = name
176        self.__flags = flags
177        self.__scope = (flags >> SCOPE_OFF) & SCOPE_MASK # like PyST_GetScope()
178        self.__namespaces = namespaces or ()
179
180    def __repr__(self):
181        return "<symbol {0!r}>".format(self.__name)
182
183    def get_name(self):
184        return self.__name
185
186    def is_referenced(self):
187        return bool(self.__flags & _symtable.USE)
188
189    def is_parameter(self):
190        return bool(self.__flags & DEF_PARAM)
191
192    def is_global(self):
193        return bool(self.__scope in (GLOBAL_IMPLICIT, GLOBAL_EXPLICIT))
194
195    def is_declared_global(self):
196        return bool(self.__scope == GLOBAL_EXPLICIT)
197
198    def is_local(self):
199        return bool(self.__flags & DEF_BOUND)
200
201    def is_free(self):
202        return bool(self.__scope == FREE)
203
204    def is_imported(self):
205        return bool(self.__flags & DEF_IMPORT)
206
207    def is_assigned(self):
208        return bool(self.__flags & DEF_LOCAL)
209
210    def is_namespace(self):
211        """Returns true if name binding introduces new namespace.
212
213        If the name is used as the target of a function or class
214        statement, this will be true.
215
216        Note that a single name can be bound to multiple objects.  If
217        is_namespace() is true, the name may also be bound to other
218        objects, like an int or list, that does not introduce a new
219        namespace.
220        """
221        return bool(self.__namespaces)
222
223    def get_namespaces(self):
224        """Return a list of namespaces bound to this name"""
225        return self.__namespaces
226
227    def get_namespace(self):
228        """Returns the single namespace bound to this name.
229
230        Raises ValueError if the name is bound to multiple namespaces.
231        """
232        if len(self.__namespaces) != 1:
233            raise ValueError, "name is bound to multiple namespaces"
234        return self.__namespaces[0]
235
236if __name__ == "__main__":
237    import os, sys
238    src = open(sys.argv[0]).read()
239    mod = symtable(src, os.path.split(sys.argv[0])[1], "exec")
240    for ident in mod.get_identifiers():
241        info = mod.lookup(ident)
242        print info, info.is_local(), info.is_namespace()
243