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)