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)
Ejemplo n.º 4
0
    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"))
Ejemplo n.º 6
0
    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)