Appease review comments

Author: Balaji Veeramani

src/sagemaker/cli/compatibility/v2/ast_transformer.py

@@ -63,6 +63,7 @@ class ASTTransformer(ast.NodeTransformer):
 
     def visit_Call(self, node):
         """Visits an ``ast.Call`` node and returns a modified node or None.
+
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
@@ -79,38 +80,41 @@ def visit_Call(self, node):
 
     def visit_Name(self, node):
         """Visits an ``ast.Name`` node and returns a modified node or None.
+
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.Name): a node that represents an identifier.
 
         Returns:
             ast.AST: if the returned node is None, the original node is removed
-            from its location. Otherwise, the original node is replaced with the
-            returned node.
+                from its location. Otherwise, the original node is replaced with
+                the returned node.
         """
         for name_checker in NAME_MODIFIERS:
             node = name_checker.check_and_modify_node(node)
         return ast.fix_missing_locations(node) if node else None
 
     def visit_Import(self, node):
         """Visits an ``ast.Import`` node and returns a modified node or None.
+
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.Import): a node that represents an import statement.
 
         Returns:
             ast.AST: if the returned node is None, the original node is removed
-            from its location. Otherwise, the original node is replaced with the
-            returned node.
+                from its location. Otherwise, the original node is replaced with
+                the returned node.
         """
         for import_checker in IMPORT_MODIFIERS:
             node = import_checker.check_and_modify_node(node)
         return ast.fix_missing_locations(node) if node else None
 
     def visit_Module(self, node):
         """Visits an ``ast.Module`` node and returns a modified node or None.
+
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         The ``ast.NodeTransformer`` walks the abstract syntax tree and modifies
@@ -121,8 +125,8 @@ def visit_Module(self, node):
 
         Returns:
             ast.AST: if the returned node is None, the original node is removed
-            from its location. Otherwise, the original node is replaced with the
-            returned node.
+                from its location. Otherwise, the original node is replaced with
+                the returned node.
         """
         self.generic_visit(node)
         for module_checker in MODULE_MODIFIERS:
@@ -131,15 +135,16 @@ def visit_Module(self, node):
 
     def visit_ImportFrom(self, node):
         """Visits an ``ast.ImportFrom`` node and returns a modified node or None.
+
         See https://docs.python.org/3/library/ast.html#ast.NodeTransformer.
 
         Args:
             node (ast.ImportFrom): a node that represents an import statement.
 
         Returns:
             ast.AST: if the returned node is None, the original node is removed
-            from its location. Otherwise, the original node is replaced with the
-            returned node.
+                from its location. Otherwise, the original node is replaced with
+                the returned node.
         """
         for import_checker in IMPORT_FROM_MODIFIERS:
             node = import_checker.check_and_modify_node(node)

src/sagemaker/cli/compatibility/v2/modifiers/serde.py

@@ -10,32 +10,31 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""Classes to modify serializer and deserializer code to be compatible with
-version 2.0 and later of the SageMaker Python SDK.
-"""
+"""Classes to modify SerDe code to be compatibile with version 2.0 and later."""
 from __future__ import absolute_import
 
 import ast
 
-import pasta
-
 from sagemaker.cli.compatibility.v2.modifiers import matching
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
 OLD_AMAZON_CLASS_NAMES = {"numpy_to_record_serializer", "record_deserializer"}
 NEW_AMAZON_CLASS_NAMES = {"RecordSerializer", "RecordDeserializer"}
+OLD_PREDICTOR_CLASS_NAMES = {
+    "_CsvSerializer",
+    "_JsonSerializer",
+    "_NpySerializer",
+    "_CsvDeserializer",
+    "BytesDeserializer",
+    "StringDeserializer",
+    "StreamDeserializer",
+    "_NumpyDeserializer",
+    "_JsonDeserializer",
+}
 
 # The values are tuples so that the object can be passed to matching.matches_any.
 OLD_CLASS_NAME_TO_NAMESPACES = {
-    "_CsvSerializer": ("sagemaker.predictor",),
-    "_JsonSerializer": ("sagemaker.predictor",),
-    "_NpySerializer": ("sagemaker.predictor",),
-    "_CsvDeserializer": ("sagemaker.predictor",),
-    "BytesDeserializer": ("sagemaker.predictor",),
-    "StringDeserializer": ("sagemaker.predictor",),
-    "StreamDeserializer": ("sagemaker.predictor",),
-    "_NumpyDeserializer": ("sagemaker.predictor",),
-    "_JsonDeserializer": ("sagemaker.predictor",),
+    class_name: ("sagemaker.predictor",) for class_name in OLD_PREDICTOR_CLASS_NAMES
 }
 OLD_CLASS_NAME_TO_NAMESPACES.update(
     {class_name: ("sagemaker.amazon.common",) for class_name in OLD_AMAZON_CLASS_NAMES}
@@ -205,7 +204,7 @@ def modify_node(self, node):
         object_name = node.id if isinstance(node, ast.Name) else node.attr
         new_class_name = OLD_OBJECT_NAME_TO_NEW_CLASS_NAME[object_name]
         namespace_name = NEW_CLASS_NAME_TO_NAMESPACES[new_class_name][0]
-        subpackage_name = namespace_name[namespace_name.find(".") + 1 :]
+        subpackage_name = namespace_name.split(".")[1]
         return ast.Call(
             func=ast.Attribute(value=ast.Name(id=subpackage_name), attr=new_class_name),
             args=[],
@@ -375,7 +374,9 @@ def __init__(self):
             for class_name in NEW_CLASS_NAMES - NEW_AMAZON_CLASS_NAMES
             if "Serializer" in class_name
         }
-        import_node = pasta.parse("from sagemaker import serializers\n").body[0]
+        import_node = ast.ImportFrom(
+            module="sagemaker", names=[ast.alias(name="serializers", asname=None)], level=0
+        )
         super().__init__(class_names, import_node)
 
 
@@ -403,5 +404,7 @@ def __init__(self):
             for class_name in NEW_CLASS_NAMES - NEW_AMAZON_CLASS_NAMES
             if "Deserializer" in class_name
         }
-        import_node = pasta.parse("from sagemaker import deserializers\n").body[0]
+        import_node = ast.ImportFrom(
+            module="sagemaker", names=[ast.alias(name="deserializers", asname=None)], level=0
+        )
         super().__init__(class_names, import_node)