def test_matches_name_or_namespaces():
name = "KMeans"
namespaces = ("sagemaker", "sagemaker.amazon.kmeans")
matches = ("KMeans()", "sagemaker.KMeans()")
for call in matches:
assert matching.matches_name_or_namespaces(ast_call(call), name, namespaces)
non_matches = ("MXNet()", "sagemaker.mxnet.MXNet()")
for call in non_matches:
assert not matching.matches_name_or_namespaces(ast_call(call), name, namespaces)
def node_should_be_modified(self, node):
"""Checks if the ``ast.Call`` node instantiates a TensorFlow estimator.
TensorFlow estimator would use``script_mode`` set. This looks for the following formats:
- ``TensorFlow``
- ``sagemaker.tensorflow.TensorFlow``
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`` is instantiating a TensorFlow estimator with ``script_mode``.
"""
is_tf_constructor = matching.matches_name_or_namespaces(
node, "TensorFlow", TF_NAMESPACES)
return is_tf_constructor and self._has_script_mode_param(node)
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 node_should_be_modified(self, node):
"""Checks if the ``ast.Call`` node calls a function of interest.
This looks for the following calls:
- ``sagemaker.get_image_uri``
- ``sagemaker.amazon_estimator.get_image_uri``
- ``get_image_uri``
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, GET_IMAGE_URI_NAME,
GET_IMAGE_URI_NAMESPACES)
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.session.ShuffleConfig``
- ``session.ShuffleConfig``
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.
"""
if isinstance(node.func, ast.Name):
return False
return matching.matches_name_or_namespaces(
node, "ShuffleConfig", ("sagemaker.session", "session"))
def node_should_be_modified(self, node):
"""Checks if the ``ast.Call`` node instantiates a TensorFlow estimator with legacy mode.
This looks for the following formats:
- ``TensorFlow``
- ``sagemaker.tensorflow.TensorFlow``
- ``sagemaker.tensorflow.estimator.TensorFlow``
Legacy mode is enabled if (1) ``script_mode`` is ``False``, ``None``, or not specified,
and (2) if ``py_version`` is ``py2`` or not specified.
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`` is instantiating a TensorFlow estimator with legacy mode.
"""
is_tf_constructor = matching.matches_name_or_namespaces(
node, "TensorFlow", TF_NAMESPACES)
return is_tf_constructor and self._is_legacy_mode(node)