address PR comment

Author: laurenyu

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

@@ -98,14 +98,14 @@ def modify_node(self, node):
         framework, is_model = _framework_from_node(node)
 
         # if framework_version is not supplied, get default and append keyword
-        framework_version = parsing.arg_value(node, FRAMEWORK_ARG)
-        if framework_version is None:
+        if matching.has_arg(node, FRAMEWORK_ARG):
+            framework_version = parsing.arg_value(node, FRAMEWORK_ARG)
+        else:
             framework_version = FRAMEWORK_DEFAULTS[framework]
             node.keywords.append(ast.keyword(arg=FRAMEWORK_ARG, value=ast.Str(s=framework_version)))
 
         # if py_version is not supplied, get a conditional default, and if not None, append keyword
-        py_version = parsing.arg_value(node, PY_ARG)
-        if py_version is None:
+        if not matching.has_arg(node, PY_ARG):
             py_version = _py_version_defaults(framework, framework_version, is_model)
             if py_version:
                 node.keywords.append(ast.keyword(arg=PY_ARG, value=ast.Str(s=py_version)))
@@ -175,13 +175,13 @@ def _version_args_needed(node, image_arg):
     Applies similar logic as ``validate_version_or_image_args``
     """
     # if image_arg is present, no need to supply version arguments
-    image_name = parsing.arg_value(node, image_arg)
-    if image_name:
+    if matching.has_arg(node, image_arg):
         return False
 
     # if framework_version is None, need args
-    framework_version = parsing.arg_value(node, FRAMEWORK_ARG)
-    if framework_version is None:
+    if matching.has_arg(node, FRAMEWORK_ARG):
+        framework_version = parsing.arg_value(node, FRAMEWORK_ARG)
+    else:
         return True
 
     # check if we expect py_version and we don't get it -- framework and model dependent

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

@@ -116,4 +116,7 @@ def has_arg(node, arg):
     Returns:
         bool: if the node has the given argument.
     """
-    return parsing.arg_value(node, arg) is not None
+    try:
+        return parsing.arg_value(node, arg) is not None
+    except KeyError:
+        return False

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

@@ -13,6 +13,8 @@
 """Functions for parsing AST nodes."""
 from __future__ import absolute_import
 
+import pasta
+
 
 def arg_from_keywords(node, arg):
     """Retrieves a keyword argument from the node's keywords.
@@ -41,10 +43,13 @@ def arg_value(node, arg):
         arg (str): the name of the argument.
 
     Returns:
-        obj: the keyword argument's value if it is present. Otherwise, this returns ``None``.
+        obj: the keyword argument's value.
+
+    Raises:
+        KeyError: if the node's keywords do not contain the argument.
     """
     keyword = arg_from_keywords(node, arg)
-    if keyword and keyword.value:
-        return getattr(keyword.value, keyword.value._fields[0], None)
+    if keyword is None:
+        raise KeyError("arg '{}' not found in call: {}".format(arg, pasta.dump(node)))
 
-    return None
+    return getattr(keyword.value, keyword.value._fields[0], None) if keyword.value else None

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

@@ -12,6 +12,8 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import pytest
+
 from sagemaker.cli.compatibility.v2.modifiers import parsing
 from tests.unit.sagemaker.cli.compatibility.v2.modifiers.ast_converter import ast_call
 
@@ -45,7 +47,14 @@ def test_arg_value():
     call = ast_call("MXNet(enable_network_isolation=True)")
     assert parsing.arg_value(call, "enable_network_isolation") is True
 
+    call = ast_call("MXNet(source_dir=None)")
+    assert parsing.arg_value(call, "source_dir") is None
+
 
 def test_arg_value_absent_keyword():
-    call = ast_call("MXNet(entry_point='run')")
-    assert parsing.arg_value(call, "framework_version") is None
+    code = "MXNet(entry_point='run')"
+
+    with pytest.raises(KeyError) as e:
+        parsing.arg_value(ast_call(code), "framework_version")
+
+    assert "arg 'framework_version' not found in call: {}".format(code) in str(e.value)