11
11
"""
12
12
13
13
import ast
14
+ from typing import Any , Dict , Optional
14
15
15
16
16
17
# Copied from orig-index
@@ -21,14 +22,14 @@ class ShortCircuitingVisitor(ast.NodeVisitor):
21
22
visiting of children is not the responsibility of the visit_ method.
22
23
"""
23
24
24
- def visit (self , node ) :
25
+ def visit (self , node : ast . AST ) -> None :
25
26
method = "visit_" + node .__class__ .__name__
26
27
visitor = getattr (self , method , self .generic_visit )
27
28
rv = visitor (node )
28
29
if rv :
29
30
self .visit_children (node )
30
31
31
- def visit_children (self , node ) :
32
+ def visit_children (self , node : ast . AST ) -> None :
32
33
for field , value in ast .iter_fields (node ):
33
34
if isinstance (value , list ):
34
35
for item in value :
@@ -37,36 +38,37 @@ def visit_children(self, node):
37
38
elif isinstance (value , ast .AST ):
38
39
self .visit (value )
39
40
40
- def generic_visit (self , node ) -> bool :
41
+ def generic_visit (self , node : ast . AST ) -> bool :
41
42
return True
42
43
43
44
44
45
class QualifiedNameSaver (ShortCircuitingVisitor ):
45
46
"""Similar to LibCST's QualifiedNameProvider except simpler and wronger"""
46
47
47
- def __init__ (self ):
48
+ def __init__ (self ) -> None :
48
49
super ().__init__ ()
49
- self .qualified_name_prefixes = {}
50
+ self .qualified_name_prefixes : Dict [ str , str ] = {}
50
51
51
52
def qualified_name (self , node : ast .AST ) -> str :
52
53
if isinstance (node , ast .Attribute ):
53
54
return self .qualified_name (node .value ) + "." + node .attr
54
55
elif isinstance (node , ast .Expr ):
55
56
return self .qualified_name (node .value )
56
57
elif isinstance (node , ast .Name ):
57
- if new := self .qualified_name_prefixes .get (node .id ):
58
+ new = self .qualified_name_prefixes .get (node .id )
59
+ if new :
58
60
return new
59
61
return f"<locals>.{ node .id } "
60
62
else :
61
63
raise ValueError (f"Complex expression: { type (node )} " )
62
64
63
- def visit_Import (self , node : ast .Import ):
65
+ def visit_Import (self , node : ast .Import ) -> None :
64
66
# .names
65
67
# alias = (identifier name, identifier? asname)
66
68
for a in node .names :
67
69
self .qualified_name_prefixes [a .asname or a .name ] = a .name
68
70
69
- def visit_ImportFrom (self , node : ast .ImportFrom ):
71
+ def visit_ImportFrom (self , node : ast .ImportFrom ) -> None :
70
72
# .identifier / .level
71
73
# .names
72
74
# alias = (identifier name, identifier? asname)
@@ -87,12 +89,12 @@ class Unknown:
87
89
88
90
89
91
class SetupFindingVisitor (QualifiedNameSaver ):
90
- def __init__ (self ):
92
+ def __init__ (self ) -> None :
91
93
super ().__init__ ()
92
- self .setup_call_args = None
93
- self .setup_call_kwargs = None
94
+ self .setup_call_args : Optional [ Dict [ str , Any ]] = None
95
+ self .setup_call_kwargs : Optional [ bool ] = None
94
96
95
- def visit_Call (self , node ) :
97
+ def visit_Call (self , node : ast . Call ) -> None :
96
98
# .func (expr, can just be name)
97
99
# .args
98
100
# .keywords
0 commit comments