Repository: aws/sagemaker-python-sdk

Issue:
bug: AttributeError: 'NoneType' object has no attribute 'module' when applying the v2 upgrade tool
**Describe the bug**

When applying the `sagemaker-upgrade-v2` automated CLI tool, earlier ast modification leads to an `AttributeError` exception.

**To reproduce**

1. Craft a script that will be modified by a SerDe import renamer
1. Apply the V2 upgrade CLI tool: 

An `AttributeError` exception is raised.

**Expected behavior**

The V2 upgrade CLI would properly upgrade the script.

**Screenshots or logs**

```
❯ cat v1.py
import sagemaker

from sagemaker.predictor import csv_serializer

csv_serializer.__doc___

❯ sagemaker-upgrade-v2 --in-file v1.py --out-file v2.py
Traceback (most recent call last):
  File "~/testvenv/bin/sagemaker-upgrade-v2", line 8, in <module>
    sys.exit(main())
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.py", line 78, in main
    _update_file(args.in_file, args.out_file)
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/sagemaker_upgrade_v2.py", line 50, in _update_file
    updater_cls(input_path=input_file, output_path=output_file).update()
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/files.py", line 72, in update
    output = self._update_ast(self._read_input_file())
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/files.py", line 86, in _update_ast
    return ASTTransformer().visit(input_ast)
  File "/usr/lib/python3.8/ast.py", line 363, in visit
    return visitor(node)
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/ast_transformer.py", line 136, in visit_Module
    self.generic_visit(node)
  File "/usr/lib/python3.8/ast.py", line 439, in generic_visit
    value = self.visit(value)
  File "/usr/lib/python3.8/ast.py", line 363, in visit
    return visitor(node)
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/ast_transformer.py", line 155, in visit_ImportFrom
    node = import_checker.check_and_modify_node(node)
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/modifiers/modifier.py", line 26, in check_and_modify_node
    if self.node_should_be_modified(node):
  File "~/testvenv/lib/python3.8/site-packages/sagemaker/cli/compatibility/v2/modifiers/image_uris.py", line 115, in node_should_be_modified
    return node.module in GET_IMAGE_URI_NAMESPACES and any(
AttributeError: 'NoneType' object has no attribute 'module'
```

**System information**
A description of your system. Please provide:
- **SageMaker Python SDK version**: 2.4.1
- **Framework name (eg. PyTorch) or algorithm (eg. KMeans)**:
- **Framework version**:
- **Python version**:
- **CPU or GPU**:
- **Custom Docker image (Y/N)**:

**Additional context**

The problem comes from [the ordering](https://github.com/aws/sagemaker-python-sdk/blob/v2.4.1/src/sagemaker/cli/compatibility/v2/ast_transformer.py#L59-L60), which the existing, isolated unit tests do not cover. The earlier renamer modifies the ast, and the later renamer cannot handle this situation:

```
 59     modifiers.serde.SerdeImportFromPredictorRenamer(),
 60     modifiers.image_uris.ImageURIRetrieveImportFromRenamer(),
```


Discussion:

Resolving patch:
diff --git a/src/sagemaker/cli/compatibility/v2/modifiers/image_uris.py b/src/sagemaker/cli/compatibility/v2/modifiers/image_uris.py
--- a/src/sagemaker/cli/compatibility/v2/modifiers/image_uris.py
+++ b/src/sagemaker/cli/compatibility/v2/modifiers/image_uris.py
@@ -112,8 +112,10 @@ def node_should_be_modified(self, node):
         Returns:
             bool: If the import statement imports ``get_image_uri`` from the correct module.
         """
-        return node.module in GET_IMAGE_URI_NAMESPACES and any(
-            name.name == GET_IMAGE_URI_NAME for name in node.names
+        return (
+            node is not None
+            and node.module in GET_IMAGE_URI_NAMESPACES
+            and any(name.name == GET_IMAGE_URI_NAME for name in node.names)
         )
 
     def modify_node(self, node):


Test patch:
diff --git a/tests/unit/sagemaker/cli/compatibility/v2/test_ast_transformer.py b/tests/unit/sagemaker/cli/compatibility/v2/test_ast_transformer.py
new file mode 100644
--- /dev/null
+++ b/tests/unit/sagemaker/cli/compatibility/v2/test_ast_transformer.py
@@ -0,0 +1,48 @@
+# 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 textwrap
+
+import pasta
+import pytest
+
+from sagemaker.cli.compatibility.v2.ast_transformer import ASTTransformer
+
+
+@pytest.fixture
+def input_code():
+    return textwrap.dedent(
+        """
+        from sagemaker.predictor import csv_serializer
+
+        csv_serializer.__doc__
+        """
+    )
+
+
+@pytest.fixture
+def output_code():
+    return textwrap.dedent(
+        """
+        from sagemaker import serializers
+
+        serializers.CSVSerializer().__doc__
+        """
+    )
+
+
+def test_simple_script(input_code, output_code):
+    input_ast = pasta.parse(input_code)
+    output_ast = ASTTransformer().visit(input_ast)
+    assert pasta.dump(output_ast) == output_code