1"""
2LocalNameDeclarations gathers name of declarations local to a node.
3LocalNodeDeclarations gathers node of declarations local to a node.
4"""
5
6from pythran.passmanager import NodeAnalysis
7
8import gast as ast
9
10
11class LocalNodeDeclarations(NodeAnalysis):
12
13    """
14    Gathers all local symbols from a function.
15
16    It should not be use from outside a function, but can be used on a function
17    (but in that case, parameters are not taken into account)
18
19    >>> import gast as ast
20    >>> from pythran import passmanager, backend
21    >>> node = ast.parse('''
22    ... def foo(a):
23    ...     b = a + 1''')
24    >>> pm = passmanager.PassManager("test")
25    >>> [name.id for name in pm.gather(LocalNodeDeclarations, node)]
26    ['b']
27    >>> node = ast.parse('''
28    ... for c in range(n):
29    ...     b = a + 1''')
30    >>> pm = passmanager.PassManager("test")
31    >>> sorted([name.id for name in pm.gather(LocalNodeDeclarations, node)])
32    ['b', 'c']
33    """
34
35    def __init__(self):
36        """ Initialize empty set as the result. """
37        self.result = set()
38        super(LocalNodeDeclarations, self).__init__()
39
40    def visit_Name(self, node):
41        """ Any node with Store context is a new declaration. """
42        if isinstance(node.ctx, ast.Store):
43            self.result.add(node)
44
45
46class LocalNameDeclarations(NodeAnalysis):
47
48    """
49    Gathers all local identifiers from a node.
50
51    >>> import gast as ast
52    >>> from pythran import passmanager, backend
53    >>> node = ast.parse('''
54    ... def foo(a):
55    ...     b = a + 1''')
56    >>> pm = passmanager.PassManager("test")
57    >>> sorted(pm.gather(LocalNameDeclarations, node))
58    ['a', 'b', 'foo']
59    """
60
61    def __init__(self):
62        """ Initialize empty set as the result. """
63        self.result = set()
64        super(LocalNameDeclarations, self).__init__()
65
66    def visit_Name(self, node):
67        """ Any node with Store or Param context is a new identifier. """
68        if isinstance(node.ctx, (ast.Store, ast.Param)):
69            self.result.add(node.id)
70
71    def visit_FunctionDef(self, node):
72        """ Function name is a possible identifier. """
73        self.result.add(node.name)
74        self.generic_visit(node)
75