def modify_node()

in src/sagemaker/cli/compatibility/v2/modifiers/image_uris.py [0:0]

• 36 lines of code
• 14 McCabe index (conditional complexity)

    def modify_node(self, node):
        """Modifies the ``ast.Call`` node to call ``image_uris.retrieve`` instead.

        And also switch the first two parameters from (region, repo) to (framework, region).

        Args:
            node (ast.Call): a node that represents a *image_uris.retrieve call.
        """
        original_args = [None] * 3
        for kw in node.keywords:
            if kw.arg == "repo_name":
                original_args[0] = ast.Str(kw.value.s)
            elif kw.arg == "repo_region":
                original_args[1] = ast.Str(kw.value.s)
            elif kw.arg == "repo_version":
                original_args[2] = ast.Str(kw.value.s)

        if len(node.args) > 0:
            original_args[1] = ast.Str(node.args[0].s)
        if len(node.args) > 1:
            original_args[0] = ast.Str(node.args[1].s)
        if len(node.args) > 2:
            original_args[2] = ast.Str(node.args[2].s)

        args = []
        for arg in original_args:
            if arg:
                args.append(arg)

        func = node.func
        has_sagemaker = False
        while hasattr(func, "value"):
            if hasattr(func.value, "id") and func.value.id == "sagemaker":
                has_sagemaker = True
                break
            func = func.value

        if has_sagemaker:
            node.func = ast.Attribute(
                value=ast.Attribute(attr="image_uris", value=ast.Name(id="sagemaker")),
                attr="retrieve",
            )
        else:
            node.func = ast.Attribute(value=ast.Name(id="image_uris"), attr="retrieve")
        node.args = args
        node.keywords = []
        return node