1"""
2Locals computes the value of locals()
3"""
4
5from pythran.passmanager import ModuleAnalysis
6import pythran.metadata as md
7
8import gast as ast
9
10
11class Locals(ModuleAnalysis):
12    """
13    Statically compute the value of locals() before each statement
14
15    Yields a dictionary binding every node to the set of variable names defined
16    *before* this node.
17
18    Following snippet illustrates its behavior:
19    >>> import gast as ast
20    >>> from pythran import passmanager
21    >>> pm = passmanager.PassManager('test')
22    >>> code = '''
23    ... def b(n):
24    ...     m = n + 1
25    ...     def b(n):
26    ...         return n + 1
27    ...     return b(m)'''
28    >>> tree = ast.parse(code)
29    >>> l = pm.gather(Locals, tree)
30    >>> sorted(l[tree.body[0].body[0]])
31    ['n']
32    >>> sorted(l[tree.body[0].body[1]])
33    ['b', 'm', 'n']
34    """
35
36    def __init__(self):
37        self.result = dict()
38        self.locals = set()
39        self.nesting = 0
40        super(Locals, self).__init__()
41
42    def generic_visit(self, node):
43        super(Locals, self).generic_visit(node)
44        if node not in self.result:
45            self.result[node] = self.result[self.expr_parent]
46
47    def store_and_visit(self, node):
48        self.expr_parent = node
49        self.result[node] = self.locals.copy()
50        self.generic_visit(node)
51
52    def visit_Module(self, node):
53        self.expr_parent = node
54        self.result[node] = self.locals
55        self.generic_visit(node)
56
57    def visit_FunctionDef(self, node):
58        # special case for nested functions
59        if self.nesting:
60            self.locals.add(node.name)
61        self.nesting += 1
62        self.expr_parent = node
63        self.result[node] = self.locals.copy()
64        parent_locals = self.locals.copy()
65        for default in node.args.defaults:
66            self.visit(default)
67        for arg in node.args.args:
68            if arg.annotation:
69                self.visit(arg.annotation)
70        if node.returns:
71            self.visit(node.returns)
72        self.locals.update(arg.id for arg in node.args.args)
73        for stmt in node.body:
74            self.visit(stmt)
75        self.locals = parent_locals
76        self.nesting -= 1
77
78    def visit_Assign(self, node):
79        self.expr_parent = node
80        self.result[node] = self.locals.copy()
81        md.visit(self, node)
82        self.visit(node.value)
83        self.locals.update(t.id for t in node.targets
84                           if isinstance(t, ast.Name))
85        for target in node.targets:
86            self.visit(target)
87
88    def visit_For(self, node):
89        self.expr_parent = node
90        self.result[node] = self.locals.copy()
91        md.visit(self, node)
92        self.visit(node.iter)
93        self.locals.add(node.target.id)
94        for stmt in node.body:
95            self.visit(stmt)
96        for stmt in node.orelse:
97            self.visit(stmt)
98
99    def visit_Import(self, node):
100        self.result[node] = self.locals.copy()
101        self.locals.update(alias.name for alias in node.names)
102
103    def visit_ImportFrom(self, node):
104        self.result[node] = self.locals.copy()
105        self.locals.update(alias.name for alias in node.names)
106
107    def visit_ExceptHandler(self, node):
108        self.expr_parent = node
109        self.result[node] = self.locals.copy()
110        if node.name:
111            self.locals.add(node.name.id)
112        node.type and self.visit(node.type)
113        for stmt in node.body:
114            self.visit(stmt)
115
116    # statements that do not define a new variable
117    visit_Return = store_and_visit
118    visit_Yield = store_and_visit
119    visit_Try = store_and_visit
120    visit_AugAssign = store_and_visit
121    visit_Print = store_and_visit
122    visit_While = store_and_visit
123    visit_If = store_and_visit
124    visit_Raise = store_and_visit
125    visit_Assert = store_and_visit
126    visit_Expr = store_and_visit
127    visit_Pass = store_and_visit
128    visit_Break = store_and_visit
129    visit_Continue = store_and_visit
130