change: add modifier for s3_input class

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

@@ -0,0 +1,97 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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 TrainingInput code to be compatible
+with version 2.0 and later of the SageMaker Python SDK.
+"""
+from __future__ import absolute_import
+
+import ast
+
+from sagemaker.cli.compatibility.v2.modifiers import matching
+from sagemaker.cli.compatibility.v2.modifiers.modifier import Modifier
+
+S3_INPUT_NAME = "s3_input"
+S3_INPUT_NAMESPACES = ("sagemaker", "sagemaker.inputs", "sagemaker.session")
+
+
+class TrainingInputConstructorRefactor(Modifier):
+    """A class to refactor *s3_input class."""
+
+    def node_should_be_modified(self, node):
+        """Checks if the ``ast.Call`` node instantiates a class of interest.
+
+        This looks for the following calls:
+
+        - ``sagemaker.s3_input``
+        - ``sagemaker.session.s3_input``
+        - ``s3_input``
+
+        Args:
+            node (ast.Call): a node that represents a function call. For more,
+                see https://docs.python.org/3/library/ast.html#abstract-grammar.
+
+        Returns:
+            bool: If the ``ast.Call`` instantiates a class of interest.
+        """
+        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.
+        """
+        if matching.matches_name(node, S3_INPUT_NAME):
+            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):
+    """A class to update import statements of ``s3_input``."""
+
+    def node_should_be_modified(self, node):
+        """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 ``s3_input`` from the correct module.
+        """
+        return node.module in S3_INPUT_NAMESPACES and any(
+            name.name == S3_INPUT_NAME for name in node.names
+        )
+
+    def modify_node(self, node):
+        """Changes the ``ast.ImportFrom`` node's name from ``s3_input`` to ``TrainingInput``.
+
+        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.
+        """
+        for name in node.names:
+            if name.name == S3_INPUT_NAME:
+                name.name = "TrainingInput"
+            if node.module == "sagemaker.session":
+                node.module = "sagemaker.inputs"

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

@@ -0,0 +1,113 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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.
+from __future__ import absolute_import
+
+import pasta
+import pytest
+
+from sagemaker.cli.compatibility.v2.modifiers import training_input
+from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call, ast_import
+
+
+@pytest.fixture
+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')",
+    )
+
+
+@pytest.fixture
+def import_statements():
+    return (
+        "from sagemaker.session import s3_input",
+        "from sagemaker.inputs import s3_input",
+        "from sagemaker import s3_input",
+    )
+
+
+def test_constructor_node_should_be_modified(constructors):
+    modifier = training_input.TrainingInputConstructorRefactor()
+    for constructor in constructors:
+        node = ast_call(constructor)
+        assert modifier.node_should_be_modified(node)
+
+
+def test_constructor_node_should_be_modified_random_call():
+    modifier = training_input.TrainingInputConstructorRefactor()
+    node = ast_call("FileSystemInput()")
+    assert not modifier.node_should_be_modified(node)
+
+
+def test_constructor_modify_node():
+    modifier = training_input.TrainingInputConstructorRefactor()
+
+    node = ast_call("s3_input(s3_data='s3://a')")
+    modifier.modify_node(node)
+    assert "TrainingInput(s3_data='s3://a')" == pasta.dump(node)
+
+    node = ast_call("sagemaker.s3_input(s3_data='s3://a')")
+    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()
+    for statement in import_statements:
+        node = ast_import(statement)
+        assert modifier.node_should_be_modified(node)
+
+
+def test_import_from_node_should_be_modified_random_import():
+    modifier = training_input.TrainingInputImportFromRenamer()
+    node = ast_import("from sagemaker.session import Session")
+    assert not modifier.node_should_be_modified(node)
+
+
+def test_import_from_modify_node():
+    modifier = training_input.TrainingInputImportFromRenamer()
+
+    node = ast_import("from sagemaker import s3_input")
+    modifier.modify_node(node)
+    expected_result = "from sagemaker import TrainingInput"
+    assert expected_result == pasta.dump(node)
+
+    node = ast_import("from sagemaker.inputs import s3_input as training_input")
+    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.inputs import TrainingInput as training_input"
+    assert expected_result == pasta.dump(node)