rename namespace

Author: Chuyang Deng

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

@@ -15,6 +15,8 @@
 """
 from __future__ import absolute_import
 
+import ast
+
 from sagemaker.cli.compatibility.v2.modifiers import matching
 from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
 
@@ -53,6 +55,15 @@ def modify_node(self, node):
             node.func.id = "TrainingInput"
         elif matching.matches_attr(node, S3_INPUT_NAME):
             node.func.attr = "TrainingInput"
+            _rename_namespace(node, "session")
+
+
+def _rename_namespace(node, name):
+    """Rename namespace ``session`` to ``inputs`` """
+    if isinstance(node.func.value, ast.Attribute) and node.func.value.attr == name:
+        node.func.value.attr = "inputs"
+    elif isinstance(node.func.value, ast.Name) and node.func.value.id == name:
+        node.func.value.id = "inputs"
 
 
 class TrainingInputImportFromRenamer(Modifier):
@@ -82,5 +93,5 @@ def modify_node(self, node):
         for name in node.names:
             if name.name == S3_INPUT_NAME:
                 name.name = "TrainingInput"
-            elif name.name == "session":
-                name.name = "inputs"
+            if node.module == "sagemaker.session":
+                node.module = "sagemaker.inputs"

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

@@ -23,7 +23,10 @@
 def constructors():
     return (
         "sagemaker.session.s3_input(s3_data='s3://a')",
+        "sagemaker.inputs.s3_input(s3_data='s3://a')",
         "sagemaker.s3_input(s3_data='s3://a')",
+        "session.s3_input(s3_data='s3://a')",
+        "inputs.s3_input(s3_data='s3://a')",
         "s3_input(s3_data='s3://a')",
     )
 
@@ -32,6 +35,7 @@ def constructors():
 def import_statements():
     return (
         "from sagemaker.session import s3_input",
+        "from sagemaker.inputs import s3_input",
         "from sagemaker import s3_input",
     )
 
@@ -60,6 +64,22 @@ def test_constructor_modify_node():
     modifier.modify_node(node)
     assert "sagemaker.TrainingInput(s3_data='s3://a')" == pasta.dump(node)
 
+    node = ast_call("session.s3_input(s3_data='s3://a')")
+    modifier.modify_node(node)
+    assert "inputs.TrainingInput(s3_data='s3://a')" == pasta.dump(node)
+
+    node = ast_call("inputs.s3_input(s3_data='s3://a')")
+    modifier.modify_node(node)
+    assert "inputs.TrainingInput(s3_data='s3://a')" == pasta.dump(node)
+
+    node = ast_call("sagemaker.inputs.s3_input(s3_data='s3://a')")
+    modifier.modify_node(node)
+    assert "sagemaker.inputs.TrainingInput(s3_data='s3://a')" == pasta.dump(node)
+
+    node = ast_call("sagemaker.session.s3_input(s3_data='s3://a')")
+    modifier.modify_node(node)
+    assert "sagemaker.inputs.TrainingInput(s3_data='s3://a')" == pasta.dump(node)
+
 
 def test_import_from_node_should_be_modified_training_input(import_statements):
     modifier = training_input.TrainingInputImportFromRenamer()
@@ -70,7 +90,7 @@ def test_import_from_node_should_be_modified_training_input(import_statements):
 
 def test_import_from_node_should_be_modified_random_import():
     modifier = training_input.TrainingInputImportFromRenamer()
-    node = ast_import("from sagemaker import Session")
+    node = ast_import("from sagemaker.session import Session")
     assert not modifier.node_should_be_modified(node)
 
 
@@ -89,5 +109,5 @@ def test_import_from_modify_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"
+    expected_result = "from sagemaker.inputs import TrainingInput as training_input"
     assert expected_result == pasta.dump(node)