sagemaker-core 1.0.47__py3-none-any.whl → 2.1.1__py3-none-any.whl

{sagemaker_core → sagemaker/core}/tools/resources_codegen.py

@@ -15,11 +15,11 @@ from functools import lru_cache
 
 import os
 import json
-from sagemaker_core.main.code_injection.codec import pascal_to_snake
-from sagemaker_core.main.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA
-from sagemaker_core.main.exceptions import DefaultConfigsError
-from sagemaker_core.main.utils import get_textual_rich_logger
-from sagemaker_core.tools.constants import (
+from sagemaker.core.utils.code_injection.codec import pascal_to_snake
+from sagemaker.core.config_schema import SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA
+from sagemaker.core.utils.exceptions import IntelligentDefaultsError
+from sagemaker.core.utils.utils import get_textual_rich_logger
+from sagemaker.core.tools.constants import (
     BASIC_RETURN_TYPES,
     GENERATED_CLASSES_LOCATION,
     RESOURCES_CODEGEN_FILE_NAME,
@@ -31,17 +31,17 @@ from sagemaker_core.tools.constants import (
     CONFIGURABLE_ATTRIBUTE_SUBSTRINGS,
     RESOURCE_WITH_LOGS,
 )
-from sagemaker_core.tools.method import Method, MethodType
-from sagemaker_core.main.utils import (
+from sagemaker.core.tools.method import Method, MethodType
+from sagemaker.core.utils.utils import (
     add_indent,
     convert_to_snake_case,
     snake_to_pascal,
     remove_html_tags,
     escape_special_rst_characters,
 )
-from sagemaker_core.tools.resources_extractor import ResourcesExtractor
-from sagemaker_core.tools.shapes_extractor import ShapesExtractor
-from sagemaker_core.tools.templates import (
+from sagemaker.core.tools.resources_extractor import ResourcesExtractor
+from sagemaker.core.tools.shapes_extractor import ShapesExtractor
+from sagemaker.core.tools.templates import (
     CALL_OPERATION_API_NO_ARG_TEMPLATE,
     CALL_OPERATION_API_TEMPLATE,
     CREATE_METHOD_TEMPLATE,
@@ -74,8 +74,10 @@ from sagemaker_core.tools.templates import (
     RESOURCE_METHOD_EXCEPTION_DOCSTRING,
     INIT_WAIT_LOGS_TEMPLATE,
     PRINT_WAIT_LOGS,
+    SERIALIZE_INPUT_ENDPOINT_TEMPLATE,
+    DESERIALIZE_RESPONSE_ENDPOINT_TEMPLATE,
 )
-from sagemaker_core.tools.data_extractor import (
+from sagemaker.core.tools.data_extractor import (
     load_combined_shapes_data,
     load_combined_operations_data,
 )
@@ -177,7 +179,7 @@ class ResourcesCodeGen:
             "import datetime",
             "import time",
             "import functools",
-            "from pydantic import validate_call, ConfigDict, BaseModel",
+            "from pydantic import validate_call",
             "from typing import Dict, List, Literal, Optional, Union, Any\n"
             "from boto3.session import Session",
             "from rich.console import Group",
@@ -186,14 +188,18 @@ class ResourcesCodeGen:
             "from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn",
             "from rich.status import Status",
             "from rich.style import Style",
-            "from sagemaker_core.main.code_injection.codec import transform",
-            "from sagemaker_core.main.code_injection.constants import Color",
-            "from sagemaker_core.main.utils import SageMakerClient, ResourceIterator, Unassigned, get_textual_rich_logger, "
+            "from sagemaker.core.shapes import *",
+            "from sagemaker.core.helper.pipeline_variable import StrPipeVar",
+            "from sagemaker.core.utils.code_injection.codec import transform",
+            "from sagemaker.core.utils.code_injection.constants import Color",
+            "from sagemaker.core.utils.utils import SageMakerClient, ResourceIterator, Unassigned, get_textual_rich_logger, "
             "snake_to_pascal, pascal_to_snake, is_not_primitive, is_not_str_dict, is_primitive_list, serialize",
-            "from sagemaker_core.main.default_configs_helper import load_default_configs_for_resource_name, get_config_value",
-            "from sagemaker_core.main.logs import MultiLogStreamHandler",
-            "from sagemaker_core.main.exceptions import *",
-            "import sagemaker_core.main.shapes as shapes",
+            "from sagemaker.core.config.config_manager import SageMakerConfig",
+            "from sagemaker.core.utils.logs import MultiLogStreamHandler",
+            "from sagemaker.core.utils.exceptions import *",
+            "from typing import ClassVar",
+            "from sagemaker.core.serializers.base import BaseSerializer",
+            "from sagemaker.core.deserializers.base import BaseDeserializer",
         ]
 
         formated_imports = "\n".join(imports)
@@ -339,11 +345,8 @@ class ResourcesCodeGen:
             str: The formatted resource class.
 
         """
-        resource_class = f"class {resource_name}(Base):\n"
-        class_documentation_string = f"Class representing resource {resource_name}\n"
-        basic_class_content = resource_class + add_indent(
-            f'"""\n{class_documentation_string}\n"""\n', 4
-        )
+        # Initialize an empty string for the resource class
+        resource_class = ""
 
         # _get_class_attributes will return value only if the resource has get or get_all method
         if class_attribute_info := self._get_class_attributes(resource_name, class_methods):
@@ -351,7 +354,11 @@ class ResourcesCodeGen:
                 class_attribute_info
             )
 
-            class_documentation_string += f"\nAttributes:\n"
+            # Start defining the class
+            resource_class = f"class {resource_name}(Base):\n"
+
+            class_documentation_string = f"Class representing resource {resource_name}\n\n"
+            class_documentation_string += f"Attributes:\n"
             class_documentation_string += self._get_shape_attr_documentation_string(
                 attributes_and_documentation
             )
@@ -418,6 +425,11 @@ class ResourcesCodeGen:
             ):
                 resource_class += add_indent(delete_method, 4)
 
+            if start_method := self._evaluate_method(
+                resource_name, "start", object_methods, resource_attributes=resource_attributes
+            ):
+                resource_class += add_indent(start_method, 4)
+
             if stop_method := self._evaluate_method(resource_name, "stop", object_methods):
                 resource_class += add_indent(stop_method, 4)
 
@@ -443,12 +455,9 @@ class ResourcesCodeGen:
         else:
             # If there's no 'get' or 'list' or 'create' method, generate a class with no attributes
             resource_attributes = []
-            resource_class = basic_class_content
-
-            # if there are no attributes or methods including addition operations
-            # class need not be generated
-            if resource_name not in self.resource_methods:
-                return ""
+            resource_class = f"class {resource_name}(Base):\n"
+            class_documentation_string = f"Class representing resource {resource_name}\n"
+            resource_class += add_indent(f'"""\n{class_documentation_string}\n"""\n', 4)
 
         if resource_name in self.resource_methods:
             # TODO: use resource_methods for all methods
@@ -496,17 +505,17 @@ class ResourcesCodeGen:
             # These resources take in the main identifier in the create and get methods , but is not present in the describe response output
             # Hence for consistent behaviour of functions such as refresh and delete, the identifiers are hardcoded
             if resource_name == "ImageVersion":
-                class_attributes["image_name"] = "str"
-                class_attributes_string = "image_name: str\n" + class_attributes_string
+                class_attributes["image_name"] = "StrPipeVar"
+                class_attributes_string = "image_name: StrPipeVar\n" + class_attributes_string
             if resource_name == "Workteam":
-                class_attributes["workteam_name"] = "str"
-                class_attributes_string = "workteam_name: str\n" + class_attributes_string
+                class_attributes["workteam_name"] = "StrPipeVar"
+                class_attributes_string = "workteam_name: StrPipeVar\n" + class_attributes_string
             if resource_name == "Workforce":
-                class_attributes["workforce_name"] = "str"
-                class_attributes_string = "workforce_name: str\n" + class_attributes_string
+                class_attributes["workforce_name"] = "StrPipeVar"
+                class_attributes_string = "workforce_name: StrPipeVar\n" + class_attributes_string
             if resource_name == "SubscribedWorkteam":
-                class_attributes["workteam_arn"] = "str"
-                class_attributes_string = "workteam_arn: str\n" + class_attributes_string
+                class_attributes["workteam_arn"] = "StrPipeVar"
+                class_attributes_string = "workteam_arn: StrPipeVar\n" + class_attributes_string
 
             if resource_name == "HubContent":
                 class_attributes["hub_name"] = "Optional[str] = Unassigned()"
@@ -514,6 +523,21 @@ class ResourcesCodeGen:
                 class_attributes_string = (
                     class_attributes_string + "hub_name: Optional[str] = Unassigned()"
                 )
+            if resource_name == "Endpoint":
+                class_attributes["serializer"] = "Optional[BaseSerializer] = None"
+                class_attributes_string = class_attributes_string.replace(
+                    "serializer: BaseSerializer", ""
+                )
+                class_attributes_string = (
+                    class_attributes_string + "serializer: Optional[BaseSerializer] = None\n"
+                )
+                class_attributes["deserializer"] = "Optional[BaseDeserializer] = None"
+                class_attributes_string = class_attributes_string.replace(
+                    "deserializer: BaseDeserializer", ""
+                )
+                class_attributes_string = (
+                    class_attributes_string + "deserializer: Optional[BaseDeserializer] = None\n"
+                )
 
             return class_attributes, class_attributes_string, attributes_and_documentation
         elif "get_all" in class_methods:
@@ -583,9 +607,11 @@ class ResourcesCodeGen:
                 and snake_to_pascal(attr[: -len("_name")]) in resource_names
             ):
                 if value.startswith("Optional"):
-                    init_data_body += f"{attr}: Optional[Union[str, object]] = Unassigned()\n"
+                    init_data_body += (
+                        f"{attr}: Optional[Union[StrPipeVar, object]] = Unassigned()\n"
+                    )
                 else:
-                    init_data_body += f"{attr}: Union[str, object]\n"
+                    init_data_body += f"{attr}: Union[StrPipeVar, object]\n"
             elif attr == "lambda":
                 init_data_body += f"# {attr}: {value}\n"
             else:
@@ -649,9 +675,9 @@ class ResourcesCodeGen:
                 and snake_to_pascal(attr[: -len("_name")]) in self.resource_names
             ):
                 if attr_type.startswith("Optional"):
-                    method_args += f"{attr}: Optional[Union[str, object]] = Unassigned(),"
+                    method_args += f"{attr}: Optional[Union[StrPipeVar, object]] = Unassigned(),"
                 else:
-                    method_args += f"{attr}: Union[str, object],"
+                    method_args += f"{attr}: Union[StrPipeVar, object],"
             else:
                 method_args += f"{attr}: {method_parameter_type},"
             if attr != last_key:
@@ -861,7 +887,7 @@ class ResourcesCodeGen:
             operation_input_shape_name=operation_input_shape_name,
             include_session_region=True,
             include_return_resource_docstring=True,
-            include_default_configs_errors=True,
+            include_intelligent_defaults_errors=True,
         )
 
         if "Describe" + resource_name in self.operations:
@@ -910,11 +936,18 @@ class ResourcesCodeGen:
             deserialize_response = DESERIALIZE_INPUT_AND_RESPONSE_TO_CLS_TEMPLATE.format(
                 operation_output_shape=operation_output_shape_name
             )
+            method_args = (
+                add_indent("cls,\n", 4)
+                + create_args
+                + "\n"
+                + add_indent("session: Optional[Session] = None,\n", 4)
+                + add_indent("region: Optional[str] = None,", 4)
+            )
             formatted_method = GENERIC_METHOD_TEMPLATE.format(
                 docstring=docstring,
                 decorator=decorator,
                 method_name="create",
-                method_args=add_indent("cls,\n", 4) + create_args,
+                method_args=method_args,
                 return_type=f'Optional["{resource_name}"]',
                 serialize_operation_input=serialize_operation_input,
                 initialize_client=initialize_client,
@@ -958,7 +991,7 @@ class ResourcesCodeGen:
         include_session_region: bool = False,
         include_return_resource_docstring: bool = False,
         return_string: str = None,
-        include_default_configs_errors: bool = False,
+        include_intelligent_defaults_errors: bool = False,
         exclude_resource_attrs: list = None,
     ) -> str:
         """
@@ -972,13 +1005,14 @@ class ResourcesCodeGen:
             include_session_region (bool): Whether to include session and region documentation.
             include_return_resource_docstring (bool): Whether to include resource-specific documentation.
             return_string (str): The return string.
-            include_default_configs_errors (bool): Whether to include default configs errors.
+            include_intelligent_defaults_errors (bool): Whether to include intelligent defaults errors.
             exclude_resource_attrs (list): A list of attributes to exclude from the docstring.
 
         Returns:
             str: The generated docstring for the IMPORT method.
         """
         docstring = f"{title}\n"
+        _shape_attr_documentation_string = ""
         if operation_input_shape_name:
             _shape_attr_documentation_string = self._get_shape_attr_documentation_string(
                 self.shapes_extractor.fetch_shape_members_and_doc_strings(
@@ -1002,8 +1036,8 @@ class ResourcesCodeGen:
 
         docstring += self._exception_docstring(operation_name)
 
-        if include_default_configs_errors:
-            subclasses = set(DefaultConfigsError.__subclasses__())
+        if include_intelligent_defaults_errors:
+            subclasses = set(IntelligentDefaultsError.__subclasses__())
             _id_exception_docstrings = [
                 f"\n    {subclass.__name__}: {subclass.__doc__}" for subclass in subclasses
             ]
@@ -1291,6 +1325,63 @@ class ResourcesCodeGen:
         )
         return formatted_method
 
+    def generate_start_method(self, resource_name: str, **kwargs) -> str:
+        """Auto-Generate 'start' object Method [delete API] for a resource.
+
+        Args:
+            resource_name (str): The resource name.
+
+        Returns:
+            str: The formatted stop Method template.
+        """
+        operation_name = "Start" + resource_name
+        operation_metadata = self.operations[operation_name]
+        operation_input_shape_name = operation_metadata["input"]["shape"]
+        resource_attributes = kwargs["resource_attributes"]
+
+        method_args = add_indent("self,\n", 4)
+        method_args += (
+            self._generate_method_args(operation_input_shape_name, resource_attributes) + "\n"
+        )
+        operation_input_args = self._generate_operation_input_args_updated(
+            operation_metadata, False, resource_attributes
+        )
+        exclude_resource_attrs = resource_attributes
+        method_args += add_indent("session: Optional[Session] = None,\n", 4)
+        method_args += add_indent("region: Optional[str] = None,", 4)
+
+        serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format(
+            operation_input_args=operation_input_args
+        )
+        call_operation_api = CALL_OPERATION_API_TEMPLATE.format(
+            operation=convert_to_snake_case(operation_name)
+        )
+
+        # generate docstring
+        docstring = self._generate_docstring(
+            title=f"Start a {resource_name} resource",
+            operation_name=operation_name,
+            resource_name=resource_name,
+            include_session_region=True,
+            include_return_resource_docstring=False,
+            exclude_resource_attrs=exclude_resource_attrs,
+        )
+
+        initialize_client = INITIALIZE_CLIENT_TEMPLATE.format(service_name="sagemaker")
+
+        formatted_method = GENERIC_METHOD_TEMPLATE.format(
+            docstring=docstring,
+            decorator="",
+            method_name="start",
+            method_args=method_args,
+            return_type="None",
+            serialize_operation_input=serialize_operation_input,
+            initialize_client=initialize_client,
+            call_operation_api=call_operation_api,
+            deserialize_response="",
+        )
+        return formatted_method
+
     def generate_stop_method(self, resource_name: str) -> str:
         """Auto-Generate 'stop' object Method [delete API] for a resource.
 
@@ -1390,12 +1481,9 @@ class ResourcesCodeGen:
                 return_type_conversion = "cls"
                 return_string = f"Returns:\n" f"    {method.resource_name}\n"
             else:
+                return_type = f"Optional[{method.return_type}]"
                 return_type_conversion = method.return_type
-                if method.resource_name != method.return_type and method.return_type in self.shapes:
-                    return_type_conversion = f"shapes.{method.return_type}"
-                return_type = f"Optional[{return_type_conversion}]"
-
-                return_string = f"Returns:\n" f"    {return_type_conversion}\n"
+                return_string = f"Returns:\n" f"    {method.return_type}\n"
             operation_output_shape = operation_metadata["output"]["shape"]
             deserialize_response = DESERIALIZE_RESPONSE_TEMPLATE.format(
                 operation_output_shape=operation_output_shape,
@@ -1405,9 +1493,21 @@ class ResourcesCodeGen:
         initialize_client = INITIALIZE_CLIENT_TEMPLATE.format(service_name=method.service_name)
         if len(self.shapes[operation_input_shape_name]["members"]) != 0:
             # the method has input arguments
-            serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format(
-                operation_input_args=operation_input_args
-            )
+            if method.resource_name == "Endpoint" and method.method_name == "invoke":
+                serialize_operation_input = SERIALIZE_INPUT_ENDPOINT_TEMPLATE.format(
+                    operation_input_args=operation_input_args
+                )
+                return_type_conversion = method.return_type
+                operation_output_shape = operation_metadata["output"]["shape"]
+                deserialize_response = DESERIALIZE_RESPONSE_ENDPOINT_TEMPLATE.format(
+                    operation_output_shape=operation_output_shape,
+                    return_type_conversion=return_type_conversion,
+                )
+
+            else:
+                serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format(
+                    operation_input_args=operation_input_args
+                )
             call_operation_api = CALL_OPERATION_API_TEMPLATE.format(
                 operation=convert_to_snake_case(method.operation_name)
             )
@@ -1476,14 +1576,10 @@ class ResourcesCodeGen:
         method_args += add_indent("session: Optional[Session] = None,\n", 4)
         method_args += add_indent("region: Optional[str] = None,", 4)
 
-        iterator_return_type = method.return_type
-        if method.resource_name != method.return_type and method.return_type in self.shapes:
-            iterator_return_type = f"shapes.{method.return_type}"
-
         if method.return_type == method.resource_name:
-            method_return_type = f'ResourceIterator["{method.resource_name}"]'
+            return_type = f'ResourceIterator["{method.resource_name}"]'
         else:
-            method_return_type = f"ResourceIterator[{iterator_return_type}]"
+            return_type = f"ResourceIterator[{method.return_type}]"
         return_string = f"Returns:\n" f"    Iterator for listed {method.return_type}.\n"
 
         get_list_operation_output_shape = operation_metadata["output"]["shape"]
@@ -1506,7 +1602,7 @@ class ResourcesCodeGen:
             f"list_method='{list_method}'",
             f"summaries_key='{summaries_key}'",
             f"summary_name='{summary_name}'",
-            f"resource_cls={iterator_return_type}",
+            f"resource_cls={method.return_type}",
             "list_method_kwargs=operation_input_args",
         ]
 
@@ -1536,7 +1632,7 @@ class ResourcesCodeGen:
             decorator=decorator,
             method_name=method.method_name,
             method_args=method_args,
-            return_type=method_return_type,
+            return_type=return_type,
             serialize_operation_input=serialize_operation_input,
             initialize_client=initialize_client,
             call_operation_api="",
@@ -1756,27 +1852,14 @@ class ResourcesCodeGen:
             get_operation_required_input = []
 
         custom_key_mapping_str = ""
-        custom_key_mapping = {}
-        extract_name_mapping_str = ""
-        extract_name_mapping = {}
         if any(member not in summary_members for member in get_operation_required_input):
-            if summary_name == "MonitoringJobDefinitionSummary":
+            if "MonitoringJobDefinitionSummary" == summary_name:
                 custom_key_mapping = {
                     "monitoring_job_definition_name": "job_definition_name",
                     "monitoring_job_definition_arn": "job_definition_arn",
                 }
-            elif summary_name == "HubContentInfo":
-                # HubContentArn -- arn:<partition>:sagemaker:<region>:<account-id>:hub-content/<hub-name>/<type>/<name>/<version>
-                # {source key from input: (target label in arn, target key in output)}
-                extract_name_mapping = {
-                    "hub_content_arn": ("hub-content/", "hub_name"),
-                }
-            elif summary_name == "ImageVersion":
-                # ImageVersionArn -- arn:aws:sagemaker:<region>:<account>:image-version/<image-name>/<version-number>
-                # {source key from input: (target label in arn, target key in output)}
-                extract_name_mapping = {
-                    "image_version_arn": ("image-version/", "image_name"),
-                }
+                custom_key_mapping_str = f"custom_key_mapping = {json.dumps(custom_key_mapping)}"
+                custom_key_mapping_str = add_indent(custom_key_mapping_str, 4)
             else:
                 log.warning(
                     f"Resource {resource_name} summaries do not have required members to create object instance. Resource may require custom key mapping for get_all().\n"
@@ -1784,15 +1867,6 @@ class ResourcesCodeGen:
                 )
                 return ""
 
-            if custom_key_mapping:
-                custom_key_mapping_str = add_indent(
-                    f"custom_key_mapping = {json.dumps(custom_key_mapping)}", 4
-                )
-            if extract_name_mapping:
-                extract_name_mapping_str = add_indent(
-                    f"extract_name_mapping = {json.dumps(extract_name_mapping)}", 4
-                )
-
         resource_iterator_args_list = [
             "client=client",
             f"list_method='{operation}'",
@@ -1804,9 +1878,6 @@ class ResourcesCodeGen:
         if custom_key_mapping_str:
             resource_iterator_args_list.append(f"custom_key_mapping=custom_key_mapping")
 
-        if extract_name_mapping_str:
-            resource_iterator_args_list.append(f"extract_name_mapping=extract_name_mapping")
-
         exclude_list = ["next_token", "max_results"]
         get_all_args = self._generate_method_args(operation_input_shape_name, exclude_list)
 
@@ -1819,7 +1890,6 @@ class ResourcesCodeGen:
                 resource=resource_name,
                 operation=operation,
                 custom_key_mapping=custom_key_mapping_str,
-                extract_name_mapping=extract_name_mapping_str,
                 resource_iterator_args=resource_iterator_args,
             )
             return formatted_method
@@ -1850,7 +1920,6 @@ class ResourcesCodeGen:
             get_all_args=get_all_args,
             operation_input_args=operation_input_args,
             custom_key_mapping=custom_key_mapping_str,
-            extract_name_mapping=extract_name_mapping_str,
             resource_iterator_args=resource_iterator_args,
         )
         return formatted_method
@@ -1937,8 +2006,6 @@ class ResourcesCodeGen:
             new_val = value.split("=")[0].strip()
             if new_val.startswith("Optional"):
                 new_val = new_val.replace("Optional[", "")[:-1]
-            if new_val.startswith("shapes."):
-                new_val = new_val.replace("shapes.", "")
             cleaned_class_attributes[key] = new_val
         return cleaned_class_attributes
 
@@ -1953,7 +2020,7 @@ class ResourcesCodeGen:
             Dict with attributes that can be configurable
 
         """
-        PYTHON_TYPES = ["str", "datetime.datetime", "bool", "int", "float"]
+        PYTHON_TYPES = ["StrPipeVar", "datetime.datetime", "bool", "int", "float"]
         default_attributes = {}
         for key, value in class_attributes.items():
             if value in PYTHON_TYPES or value.startswith("List"):

{sagemaker_core → sagemaker/core}/tools/resources_extractor.py

@@ -15,14 +15,14 @@ from typing import Optional
 
 import pandas as pd
 
-from sagemaker_core.main.utils import get_textual_rich_logger
-from sagemaker_core.tools.constants import CLASS_METHODS, OBJECT_METHODS
-from sagemaker_core.tools.data_extractor import (
+from sagemaker.core.utils.utils import get_textual_rich_logger
+from sagemaker.core.tools.constants import CLASS_METHODS, OBJECT_METHODS
+from sagemaker.core.tools.data_extractor import (
     load_additional_operations_data,
     load_combined_operations_data,
     load_combined_shapes_data,
 )
-from sagemaker_core.tools.method import Method
+from sagemaker.core.tools.method import Method
 
 log = get_textual_rich_logger(__name__)
 """
@@ -311,15 +311,7 @@ class ResourcesExtractor:
                             ):
                                 chain_resource_names.add(chain_resource_name)
                 action_split = action_low.split(resource_low)
-                if action_split[0] == "invoke":
-                    if not action_split[1]:
-                        invoke_method = "invoke"
-                    elif action_split[1] == "async":
-                        invoke_method = "invoke_async"
-                    else:
-                        invoke_method = "invoke_with_response_stream"
-                    object_methods.add(invoke_method)
-                elif action_split[0] in CLASS_METHODS:
+                if action_split[0] in CLASS_METHODS:
                     if action_low.split(resource_low)[0] == "list":
                         class_methods.add("get_all")
                     else:

{sagemaker_core → sagemaker/core}/tools/shapes_codegen.py

@@ -17,22 +17,21 @@ export PYTHONPATH=<sagemaker-code-gen repo directory>:$PYTHONPATH
 """
 import os
 
-from sagemaker_core.main.code_injection.codec import pascal_to_snake
-from sagemaker_core.tools.constants import (
+from sagemaker.core.utils.code_injection.codec import pascal_to_snake
+from sagemaker.core.tools.constants import (
     LICENCES_STRING,
     GENERATED_CLASSES_LOCATION,
     SHAPES_CODEGEN_FILE_NAME,
-    SHAPES_WITH_JSON_FIELD_ALIAS,
 )
-from sagemaker_core.tools.shapes_extractor import ShapesExtractor
-from sagemaker_core.main.utils import (
+from sagemaker.core.tools.shapes_extractor import ShapesExtractor
+from sagemaker.core.utils.utils import (
     add_indent,
     convert_to_snake_case,
     remove_html_tags,
     escape_special_rst_characters,
 )
-from sagemaker_core.tools.templates import SHAPE_CLASS_TEMPLATE, SHAPE_BASE_CLASS_TEMPLATE
-from sagemaker_core.tools.data_extractor import (
+from sagemaker.core.tools.templates import SHAPE_CLASS_TEMPLATE, SHAPE_BASE_CLASS_TEMPLATE
+from sagemaker.core.tools.data_extractor import (
     load_combined_shapes_data,
     load_combined_operations_data,
 )
@@ -149,8 +148,9 @@ class ShapesCodeGen:
         """
         class_name = shape
         init_data = self.shapes_extractor.generate_data_shape_string_body(
-            shape, self.resources_plan, add_shapes_prefix=False
+            shape, self.resources_plan
         )
+
         try:
             data_class_members = add_indent(init_data, 4)
         except Exception:
@@ -181,13 +181,7 @@ class ShapesCodeGen:
 
         if "members" in shape_dict:
             for member, member_attributes in shape_dict["members"].items():
-                # Add alias if field name is json, to address the Bug: https://github.com/aws/sagemaker-python-sdk/issues/4944
-                if shape in SHAPES_WITH_JSON_FIELD_ALIAS and member == "Json":
-                    updated_member = "JsonFormat"
-                    docstring += f"\n{convert_to_snake_case(updated_member)}"
-                else:
-                    docstring += f"\n{convert_to_snake_case(member)}"
-
+                docstring += f"\n{convert_to_snake_case(member)}"
                 if "documentation" in member_attributes:
                     docstring += f": {member_attributes['documentation']}"
 
@@ -210,10 +204,15 @@ class ShapesCodeGen:
         :return: The generated import statements as a string.
         """
         imports = "import datetime\n"
+        imports += "import warnings\n"
         imports += "\n"
-        imports += "from pydantic import BaseModel, ConfigDict, Field\n"
+        imports += "from pydantic import BaseModel, ConfigDict\n"
         imports += "from typing import List, Dict, Optional, Any, Union\n"
-        imports += "from sagemaker_core.main.utils import Unassigned"
+        imports += "from sagemaker.core.utils.utils import Unassigned\n"
+        imports += "from sagemaker.core.helper.pipeline_variable import StrPipeVar\n"
+        imports += "\n"
+        imports += "# Suppress Pydantic warnings about field names shadowing parent attributes\n"
+        imports += "warnings.filterwarnings('ignore', message='.*shadows an attribute.*')\n"
         imports += "\n"
         return imports