29
29
30
30
31
31
class Visitor (ast .NodeVisitor ):
32
- def __init__ (self ):
32
+ def __init__ (self ) -> None :
33
33
self .pandas_namespace : MutableMapping [Offset , str ] = {}
34
34
self .no_namespace : Set [str ] = set ()
35
35
@@ -42,13 +42,15 @@ def visit_Attribute(self, node: ast.Attribute) -> None:
42
42
self .pandas_namespace [Offset (node .lineno , node .col_offset )] = node .attr
43
43
self .generic_visit (node )
44
44
45
- def visit_Name (self , node : ast .Name ):
45
+ def visit_Name (self , node : ast .Name ) -> None :
46
46
if node .id not in EXCLUDE :
47
47
self .no_namespace .add (node .id )
48
48
self .generic_visit (node )
49
49
50
50
51
- def check_for_inconsistent_pandas_namespace (content , path , * , replace ):
51
+ def check_for_inconsistent_pandas_namespace (
52
+ content : str , path : str , * , replace : bool
53
+ ) -> Optional [str ]:
52
54
tree = ast .parse (content )
53
55
54
56
visitor = Visitor ()
@@ -76,7 +78,8 @@ def check_for_inconsistent_pandas_namespace(content, path, *, replace):
76
78
# Replace `.`
77
79
tokens [n + 1 ] = i ._replace (src = "" )
78
80
79
- return tokens_to_src (tokens )
81
+ new_src : str = tokens_to_src (tokens )
82
+ return new_src
80
83
81
84
82
85
def main (argv : Optional [Sequence [str ]] = None ) -> None :
@@ -91,7 +94,7 @@ def main(argv: Optional[Sequence[str]] = None) -> None:
91
94
new_content = check_for_inconsistent_pandas_namespace (
92
95
content , path , replace = args .replace
93
96
)
94
- if not args .replace :
97
+ if not args .replace or new_content is None :
95
98
continue
96
99
with open (path , "w" , encoding = "utf-8" ) as fd :
97
100
fd .write (new_content )
0 commit comments