3535 modifiers .renamed_params .SessionCreateEndpointImageURIRenamer (),
3636 modifiers .training_params .TrainPrefixRemover (),
3737 modifiers .training_input .TrainingInputConstructorRefactor (),
38+ modifiers .serde .SerdeConstructorRenamer (),
3839]
3940
4041IMPORT_MODIFIERS = [modifiers .tfs .TensorFlowServingImportRenamer ()]
4142
43+ NAME_MODIFIERS = [modifiers .serde .SerdeObjectRenamer ()]
44+
45+ MODULE_MODIFIERS = [
46+ modifiers .serde .SerializerImportInserter (),
47+ modifiers .serde .DeserializerImportInserter (),
48+ ]
49+
4250IMPORT_FROM_MODIFIERS = [
4351 modifiers .predictors .PredictorImportFromRenamer (),
4452 modifiers .tfs .TensorFlowServingImportFromRenamer (),
4553 modifiers .training_input .TrainingInputImportFromRenamer (),
54+ modifiers .serde .SerdeImportFromAmazonCommonRenamer (),
55+ modifiers .serde .SerdeImportFromPredictorRenamer (),
4656]
4757
4858
@@ -52,52 +62,90 @@ class ASTTransformer(ast.NodeTransformer):
5262 """
5363
5464 def visit_Call (self , node ):
55- """Visits an ``ast.Call`` node and returns a modified node, if needed.
65+ """Visits an ``ast.Call`` node and returns a modified node or None.
66+
5667 See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
5768
5869 Args:
5970 node (ast.Call): a node that represents a function call.
6071
6172 Returns:
62- ast.Call: a node that represents a function call, which has
63- potentially been modified from the original input.
73+ ast.AST: if the returned node is None, the original node is removed
74+ from its location. Otherwise, the original node is replaced with
75+ the returned node.
6476 """
6577 for function_checker in FUNCTION_CALL_MODIFIERS :
66- function_checker .check_and_modify_node (node )
78+ node = function_checker .check_and_modify_node (node )
79+ return ast .fix_missing_locations (node ) if node else None
80+
81+ def visit_Name (self , node ):
82+ """Visits an ``ast.Name`` node and returns a modified node or None.
6783
68- ast .fix_missing_locations (node )
69- return node
84+ See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
85+
86+ Args:
87+ node (ast.Name): a node that represents an identifier.
88+
89+ Returns:
90+ ast.AST: if the returned node is None, the original node is removed
91+ from its location. Otherwise, the original node is replaced with
92+ the returned node.
93+ """
94+ for name_checker in NAME_MODIFIERS :
95+ node = name_checker .check_and_modify_node (node )
96+ return ast .fix_missing_locations (node ) if node else None
7097
7198 def visit_Import (self , node ):
72- """Visits an ``ast.Import`` node and returns a modified node, if needed.
99+ """Visits an ``ast.Import`` node and returns a modified node or None.
100+
73101 See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
74102
75103 Args:
76104 node (ast.Import): a node that represents an import statement.
77105
78106 Returns:
79- ast.Import: a node that represents an import statement, which has
80- potentially been modified from the original input.
107+ ast.AST: if the returned node is None, the original node is removed
108+ from its location. Otherwise, the original node is replaced with
109+ the returned node.
81110 """
82111 for import_checker in IMPORT_MODIFIERS :
83- import_checker .check_and_modify_node (node )
112+ node = import_checker .check_and_modify_node (node )
113+ return ast .fix_missing_locations (node ) if node else None
114+
115+ def visit_Module (self , node ):
116+ """Visits an ``ast.Module`` node and returns a modified node or None.
117+
118+ See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
119+
120+ The ``ast.NodeTransformer`` walks the abstract syntax tree and modifies
121+ all other nodes before modifying the ``ast.Module`` node.
84122
85- ast .fix_missing_locations (node )
86- return node
123+ Args:
124+ node (ast.Module): a node that represents a Python module.
125+
126+ Returns:
127+ ast.AST: if the returned node is None, the original node is removed
128+ from its location. Otherwise, the original node is replaced with
129+ the returned node.
130+ """
131+ self .generic_visit (node )
132+ for module_checker in MODULE_MODIFIERS :
133+ node = module_checker .check_and_modify_node (node )
134+ return ast .fix_missing_locations (node ) if node else None
87135
88136 def visit_ImportFrom (self , node ):
89- """Visits an ``ast.ImportFrom`` node and returns a modified node, if needed.
137+ """Visits an ``ast.ImportFrom`` node and returns a modified node or None.
138+
90139 See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
91140
92141 Args:
93142 node (ast.ImportFrom): a node that represents an import statement.
94143
95144 Returns:
96- ast.ImportFrom: a node that represents an import statement, which has
97- potentially been modified from the original input.
145+ ast.AST: if the returned node is None, the original node is removed
146+ from its location. Otherwise, the original node is replaced with
147+ the returned node.
98148 """
99149 for import_checker in IMPORT_FROM_MODIFIERS :
100- import_checker .check_and_modify_node (node )
101-
102- ast .fix_missing_locations (node )
103- return node
150+ node = import_checker .check_and_modify_node (node )
151+ return ast .fix_missing_locations (node ) if node else None
0 commit comments