10a8c90248264a8b26970b4473770bcc3df8515fJosh Gao"""Module symbol-table generator"""
20a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
30a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom compiler import ast
40a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom compiler.consts import SC_LOCAL, SC_GLOBAL_IMPLICIT, SC_GLOBAL_EXPLICIT, \
50a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    SC_FREE, SC_CELL, SC_UNKNOWN
60a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom compiler.misc import mangle
70a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport types
80a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
90a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
100a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport sys
110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
120a8c90248264a8b26970b4473770bcc3df8515fJosh GaoMANGLE_LEN = 256
130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
140a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass Scope:
150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # XXX how much information do I need about each name?
160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self, name, module, klass=None):
170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.name = name
180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.module = module
190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.defs = {}
200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.uses = {}
210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.globals = {}
220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.params = {}
230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.frees = {}
240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.cells = {}
250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.children = []
260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # nested is true if the class could contain free variables,
270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # i.e. if it is nested within another function.
280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.nested = None
290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.generator = None
300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.klass = None
310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if klass is not None:
320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            for i in range(len(klass)):
330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if klass[i] != '_':
340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    self.klass = klass[i:]
350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    break
360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __repr__(self):
380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return "<%s: %s>" % (self.__class__.__name__, self.name)
390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def mangle(self, name):
410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if self.klass is None:
420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return name
430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return mangle(name, self.klass)
440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_def(self, name):
460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.defs[self.mangle(name)] = 1
470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_use(self, name):
490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.uses[self.mangle(name)] = 1
500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_global(self, name):
520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        name = self.mangle(name)
530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.uses or name in self.defs:
540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            pass # XXX warn about global following def/use
550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.params:
560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            raise SyntaxError, "%s in %s is global and parameter" % \
570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                  (name, self.name)
580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.globals[name] = 1
590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.module.add_def(name)
600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_param(self, name):
620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        name = self.mangle(name)
630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.defs[name] = 1
640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.params[name] = 1
650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_names(self):
670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d = {}
680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d.update(self.defs)
690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d.update(self.uses)
700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d.update(self.globals)
710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return d.keys()
720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_child(self, child):
740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.children.append(child)
750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_children(self):
770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return self.children
780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def DEBUG(self):
800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, self.name, self.nested and "nested" or ""
810a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, "\tglobals: ", self.globals
820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, "\tcells: ", self.cells
830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, "\tdefs: ", self.defs
840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, "\tuses: ", self.uses
850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print >> sys.stderr, "\tfrees:", self.frees
860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def check_name(self, name):
880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """Return scope of name.
890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        The scope of a name could be LOCAL, GLOBAL, FREE, or CELL.
910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """
920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.globals:
930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_GLOBAL_EXPLICIT
940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.cells:
950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_CELL
960a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.defs:
970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_LOCAL
980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if self.nested and (name in self.frees or name in self.uses):
990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_FREE
1000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if self.nested:
1010a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_UNKNOWN
1020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        else:
1030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return SC_GLOBAL_IMPLICIT
1040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_free_vars(self):
1060a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if not self.nested:
1070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return ()
1080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        free = {}
1090a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        free.update(self.frees)
1100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name in self.uses.keys():
1110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if name not in self.defs and name not in self.globals:
1120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                free[name] = 1
1130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return free.keys()
1140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def handle_children(self):
1160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for child in self.children:
1170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            frees = child.get_free_vars()
1180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            globals = self.add_frees(frees)
1190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            for name in globals:
1200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                child.force_global(name)
1210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def force_global(self, name):
1230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """Force name to be global in scope.
1240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        Some child of the current node had a free reference to name.
1260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        When the child was processed, it was labelled a free
1270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        variable.  Now that all its enclosing scope have been
1280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        processed, the name is known to be a global or builtin.  So
1290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        walk back down the child chain and set the name to be global
1300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        rather than free.
1310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        Be careful to stop if a child does not think the name is
1330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        free.
1340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """
1350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.globals[name] = 1
1360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if name in self.frees:
1370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            del self.frees[name]
1380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for child in self.children:
1390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if child.check_name(name) == SC_FREE:
1400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                child.force_global(name)
1410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def add_frees(self, names):
1430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """Process list of free vars from nested scope.
1440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        Returns a list of names that are either 1) declared global in the
1460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        parent or 2) undefined in a top-level parent.  In either case,
1470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        the nested scope should treat them as globals.
1480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """
1490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        child_globals = []
1500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name in names:
1510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            sc = self.check_name(name)
1520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if self.nested:
1530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if sc == SC_UNKNOWN or sc == SC_FREE \
1540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                   or isinstance(self, ClassScope):
1550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    self.frees[name] = 1
1560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                elif sc == SC_GLOBAL_IMPLICIT:
1570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    child_globals.append(name)
1580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                elif isinstance(self, FunctionScope) and sc == SC_LOCAL:
1590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    self.cells[name] = 1
1600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                elif sc != SC_CELL:
1610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    child_globals.append(name)
1620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            else:
1630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if sc == SC_LOCAL:
1640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    self.cells[name] = 1
1650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                elif sc != SC_CELL:
1660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    child_globals.append(name)
1670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return child_globals
1680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_cell_vars(self):
1700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return self.cells.keys()
1710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1720a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass ModuleScope(Scope):
1730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __super_init = Scope.__init__
1740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self):
1760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__super_init("global", self)
1770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1780a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass FunctionScope(Scope):
1790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    pass
1800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1810a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass GenExprScope(Scope):
1820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __super_init = Scope.__init__
1830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __counter = 1
1850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self, module, klass=None):
1870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        i = self.__counter
1880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__counter += 1
1890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__super_init("generator expression<%d>"%i, module, klass)
1900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.add_param('.0')
1910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_names(self):
1930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        keys = Scope.get_names(self)
1940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return keys
1950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1960a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass LambdaScope(FunctionScope):
1970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __super_init = Scope.__init__
1980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
1990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __counter = 1
2000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2010a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self, module, klass=None):
2020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        i = self.__counter
2030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__counter += 1
2040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__super_init("lambda.%d" % i, module, klass)
2050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2060a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass ClassScope(Scope):
2070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    __super_init = Scope.__init__
2080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2090a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self, name, module):
2100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.__super_init(name, module, name)
2110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2120a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass SymbolVisitor:
2130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def __init__(self):
2140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.scopes = {}
2150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.klass = None
2160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # node that define new scopes
2180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitModule(self, node):
2200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope = self.module = self.scopes[node] = ModuleScope()
2210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.node, scope)
2220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    visitExpression = visitModule
2240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitFunction(self, node, parent):
2260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.decorators:
2270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.decorators, parent)
2280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        parent.add_def(node.name)
2290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for n in node.defaults:
2300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(n, parent)
2310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope = FunctionScope(node.name, self.module, self.klass)
2320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if parent.nested or isinstance(parent, FunctionScope):
2330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.nested = 1
2340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.scopes[node] = scope
2350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._do_args(scope, node.argnames)
2360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.code, scope)
2370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.handle_free_vars(scope, parent)
2380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitGenExpr(self, node, parent):
2400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope = GenExprScope(self.module, self.klass);
2410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if parent.nested or isinstance(parent, FunctionScope) \
2420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                or isinstance(parent, GenExprScope):
2430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.nested = 1
2440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.scopes[node] = scope
2460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.code, scope)
2470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.handle_free_vars(scope, parent)
2490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitGenExprInner(self, node, scope):
2510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for genfor in node.quals:
2520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(genfor, scope)
2530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope)
2550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitGenExprFor(self, node, scope):
2570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.assign, scope, 1)
2580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.iter, scope)
2590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for if_ in node.ifs:
2600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(if_, scope)
2610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitGenExprIf(self, node, scope):
2630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.test, scope)
2640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitLambda(self, node, parent, assign=0):
2660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # Lambda is an expression, so it could appear in an expression
2670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # context where assign is passed.  The transformer should catch
2680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # any code that has a lambda on the left-hand side.
2690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        assert not assign
2700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for n in node.defaults:
2720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(n, parent)
2730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope = LambdaScope(self.module, self.klass)
2740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if parent.nested or isinstance(parent, FunctionScope):
2750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.nested = 1
2760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.scopes[node] = scope
2770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self._do_args(scope, node.argnames)
2780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.code, scope)
2790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.handle_free_vars(scope, parent)
2800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2810a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def _do_args(self, scope, args):
2820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name in args:
2830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if type(name) == types.TupleType:
2840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                self._do_args(scope, name)
2850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            else:
2860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                scope.add_param(name)
2870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def handle_free_vars(self, scope, parent):
2890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        parent.add_child(scope)
2900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope.handle_children()
2910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
2920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitClass(self, node, parent):
2930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        parent.add_def(node.name)
2940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for n in node.bases:
2950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(n, parent)
2960a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope = ClassScope(node.name, self.module)
2970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if parent.nested or isinstance(parent, FunctionScope):
2980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.nested = 1
2990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.doc is not None:
3000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_def('__doc__')
3010a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope.add_def('__module__')
3020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.scopes[node] = scope
3030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        prev = self.klass
3040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.klass = node.name
3050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.code, scope)
3060a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.klass = prev
3070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.handle_free_vars(scope, parent)
3080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3090a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # name can be a def or a use
3100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # XXX a few calls and nodes expect a third "assign" arg that is
3120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # true if the name is being used as an assignment.  only
3130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # expressions contained within statements may have the assign arg.
3140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitName(self, node, scope, assign=0):
3160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if assign:
3170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_def(node.name)
3180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        else:
3190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_use(node.name)
3200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # operations that bind new names
3220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitFor(self, node, scope):
3240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.assign, scope, 1)
3250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.list, scope)
3260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.body, scope)
3270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.else_:
3280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.else_, scope)
3290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitFrom(self, node, scope):
3310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name, asname in node.names:
3320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if name == "*":
3330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                continue
3340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_def(asname or name)
3350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitImport(self, node, scope):
3370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name, asname in node.names:
3380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            i = name.find(".")
3390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if i > -1:
3400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                name = name[:i]
3410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_def(asname or name)
3420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitGlobal(self, node, scope):
3440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for name in node.names:
3450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            scope.add_global(name)
3460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitAssign(self, node, scope):
3480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """Propagate assignment flag down to child nodes.
3490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        The Assign node doesn't itself contains the variables being
3510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        assigned to.  Instead, the children in node.nodes are visited
3520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        with the assign flag set to true.  When the names occur in
3530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        those nodes, they are marked as defs.
3540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        Some names that occur in an assignment target are not bound by
3560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        the assignment, e.g. a name occurring inside a slice.  The
3570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        visitor handles these nodes specially; they do not propagate
3580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        the assign flag to their children.
3590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """
3600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for n in node.nodes:
3610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(n, scope, 1)
3620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope)
3630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitAssName(self, node, scope, assign=1):
3650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope.add_def(node.name)
3660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitAssAttr(self, node, scope, assign=0):
3680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope, 0)
3690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitSubscript(self, node, scope, assign=0):
3710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope, 0)
3720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for n in node.subs:
3730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(n, scope, 0)
3740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3750a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitSlice(self, node, scope, assign=0):
3760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope, 0)
3770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.lower:
3780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.lower, scope, 0)
3790a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.upper:
3800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.upper, scope, 0)
3810a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3820a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitAugAssign(self, node, scope):
3830a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # If the LHS is a name, then this counts as assignment.
3840a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # Otherwise, it's just use.
3850a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.node, scope)
3860a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if isinstance(node.node, ast.Name):
3870a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.node, scope, 1) # XXX worry about this
3880a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.expr, scope)
3890a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3900a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # prune if statements if tests are false
3910a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3920a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    _const_types = types.StringType, types.IntType, types.FloatType
3930a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
3940a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitIf(self, node, scope):
3950a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for test, body in node.tests:
3960a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if isinstance(test, ast.Const):
3970a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if type(test.value) in self._const_types:
3980a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    if not test.value:
3990a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        continue
4000a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(test, scope)
4010a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(body, scope)
4020a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if node.else_:
4030a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            self.visit(node.else_, scope)
4040a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4050a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    # a yield statement signals a generator
4060a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4070a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def visitYield(self, node, scope):
4080a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scope.generator = 1
4090a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.visit(node.value, scope)
4100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4110a8c90248264a8b26970b4473770bcc3df8515fJosh Gaodef list_eq(l1, l2):
4120a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    return sorted(l1) == sorted(l2)
4130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4140a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoif __name__ == "__main__":
4150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    import sys
4160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    from compiler import parseFile, walk
4170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    import symtable
4180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def get_names(syms):
4200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return [s for s in [s.get_name() for s in syms.get_symbols()]
4210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if not (s.startswith('_[') or s.startswith('.'))]
4220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    for file in sys.argv[1:]:
4240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        print file
4250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        f = open(file)
4260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        buf = f.read()
4270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        f.close()
4280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        syms = symtable.symtable(buf, file, "exec")
4290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        mod_names = get_names(syms)
4300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        tree = parseFile(file)
4310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        s = SymbolVisitor()
4320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        walk(tree, s)
4330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # compare module-level symbols
4350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        names2 = s.scopes[tree].get_names()
4360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if not list_eq(mod_names, names2):
4380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            print
4390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            print "oops", file
4400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            print sorted(mod_names)
4410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            print sorted(names2)
4420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            sys.exit(-1)
4430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d = {}
4450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        d.update(s.scopes)
4460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del d[tree]
4470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        scopes = d.values()
4480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        del d
4490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
4500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for s in syms.get_symbols():
4510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            if s.is_namespace():
4520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                l = [sc for sc in scopes
4530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                     if sc.name == s.get_name()]
4540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                if len(l) > 1:
4550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    print "skipping", s.get_name()
4560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                else:
4570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                    if not list_eq(get_names(s.get_namespace()),
4580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                                   l[0].get_names()):
4590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        print s.get_name()
4600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        print sorted(get_names(s.get_namespace()))
4610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        print sorted(l[0].get_names())
4620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        sys.exit(-1)
463