modify namespaces

Author: Chuyang Deng

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

@@ -34,13 +34,15 @@
     modifiers.renamed_params.SessionCreateModelImageURIRenamer(),
     modifiers.renamed_params.SessionCreateEndpointImageURIRenamer(),
     modifiers.training_params.TrainPrefixRemover(),
+    modifiers.training_input.TrainingInputConstructorRefactor(),
 ]
 
 IMPORT_MODIFIERS = [modifiers.tfs.TensorFlowServingImportRenamer()]
 
 IMPORT_FROM_MODIFIERS = [
     modifiers.predictors.PredictorImportFromRenamer(),
     modifiers.tfs.TensorFlowServingImportFromRenamer(),
+    modifiers.training_input.TrainingInputImportFromRenamer(),
 ]
 
 

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

@@ -22,4 +22,5 @@
     tf_legacy_mode,
     tfs,
     training_params,
+    training_input,
 )

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

@@ -10,17 +10,16 @@
 # 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 Predictor code to be compatible
+"""Classes to modify TrainingInput code to be compatible
 with version 2.0 and later of the SageMaker Python SDK.
 """
 from __future__ import absolute_import
 
 from sagemaker.cli.compatibility.v2.modifiers import matching
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
-BASE_S3_INPUT = "s3_input"
-SESSION = "session"
-S3_INPUT = {"s3_input": ("sagemaker", "sagemaker.session")}
+S3_INPUT_NAME = "s3_input"
+S3_INPUT_NAMESPACES = ("sagemaker", "sagemaker.inputs", "sagemaker.session")
 
 
 class TrainingInputConstructorRefactor(Modifier):
@@ -42,40 +41,35 @@ def node_should_be_modified(self, node):
         Returns:
             bool: If the ``ast.Call`` instantiates a class of interest.
         """
-        return matching.matches_any(node, S3_INPUT)
+        return matching.matches_name_or_namespaces(node, S3_INPUT_NAME, S3_INPUT_NAMESPACES)
 
     def modify_node(self, node):
         """Modifies the ``ast.Call`` node to call ``TrainingInput`` instead.
 
         Args:
             node (ast.Call): a node that represents a *TrainingInput constructor.
         """
-        _rename_class(node)
-
-
-def _rename_class(node):
-    """Renames the s3_input class to TrainingInput"""
-    if matching.matches_name(node, BASE_S3_INPUT):
-        node.func.id = "TrainingInput"
-    elif matching.matches_attr(node, BASE_S3_INPUT):
-        node.func.attr = "TrainingInput"
+        if matching.matches_name(node, S3_INPUT_NAME):
+            node.func.id = "TrainingInput"
+        elif matching.matches_attr(node, S3_INPUT_NAME):
+            node.func.attr = "TrainingInput"
 
 
 class TrainingInputImportFromRenamer(Modifier):
     """A class to update import statements of ``s3_input``."""
 
     def node_should_be_modified(self, node):
-        """Checks if the import statement imports ``RealTimePredictor`` from the correct module.
+        """Checks if the import statement imports ``s3_input`` from the correct module.
 
         Args:
             node (ast.ImportFrom): a node that represents a ``from ... import ... `` statement.
                 For more, see https://docs.python.org/3/library/ast.html#abstract-grammar.
 
         Returns:
-            bool: If the import statement imports ``RealTimePredictor`` from the correct module.
+            bool: If the import statement imports ``s3_input`` from the correct module.
         """
-        return node.module in S3_INPUT[BASE_S3_INPUT] and any(
-            name.name == BASE_S3_INPUT for name in node.names
+        return node.module in S3_INPUT_NAMESPACES and any(
+            name.name == S3_INPUT_NAME for name in node.names
         )
 
     def modify_node(self, node):
@@ -86,7 +80,7 @@ def modify_node(self, node):
                 For more, see https://docs.python.org/3/library/ast.html#abstract-grammar.
         """
         for name in node.names:
-            if name.name == BASE_S3_INPUT:
+            if name.name == S3_INPUT_NAME:
                 name.name = "TrainingInput"
             elif name.name == "session":
                 name.name = "inputs"

tests/unit/sagemaker/cli/compatibility/v2/modifiers/test_training_input.py

@@ -86,3 +86,8 @@ def test_import_from_modify_node():
     modifier.modify_node(node)
     expected_result = "from sagemaker.inputs import TrainingInput as training_input"
     assert expected_result == pasta.dump(node)
+
+    node = ast_import("from sagemaker.session import s3_input as training_input")
+    modifier.modify_node(node)
+    expected_result = "from sagemaker.session import TrainingInput as training_input"
+    assert expected_result == pasta.dump(node)