to_json to inherit from JumpStartDataHolderType

Author: jinyoung-lim

src/sagemaker/jumpstart/curated_hub/types.py

@@ -153,30 +153,12 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hub_content_version: str = json_obj["hub_content_version"]
         self.hub_name: str = json_obj["hub_name"]
 
-    def to_json(self) -> Dict[str, Any]:
-        """Returns json representation of DescribeHubContentsResponse object."""
-        json_obj = {}
-        for att in self.__slots__:
-            if hasattr(self, att):
-                cur_val = getattr(self, att)
-                if issubclass(type(cur_val), JumpStartDataHolderType):
-                    json_obj[att] = cur_val.to_json()
-                elif isinstance(cur_val, list):
-                    json_obj[att] = []
-                    for obj in cur_val:
-                        if issubclass(type(obj), JumpStartDataHolderType):
-                            json_obj[att].append(obj.to_json())
-                        else:
-                            json_obj[att].append(obj)
-                else:
-                    json_obj[att] = cur_val
-        return json_obj
-
 
 class HubS3StorageConfig(JumpStartDataHolderType):
     """Data class for any dependencies related to hub content.
 
-    Includes scripts, model artifacts, datasets, or notebooks."""
+    Includes scripts, model artifacts, datasets, or notebooks.
+    """
 
     __slots__ = ["s3_output_path"]
 
@@ -245,22 +227,3 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.s3_storage_config: HubS3StorageConfig = HubS3StorageConfig(
             json_obj["s3_storage_config"]
         )
-
-    def to_json(self) -> Dict[str, Any]:
-        """Returns json representation of DescribeHubContentsResponse object."""
-        json_obj = {}
-        for att in self.__slots__:
-            if hasattr(self, att):
-                cur_val = getattr(self, att)
-                if issubclass(type(cur_val), JumpStartDataHolderType):
-                    json_obj[att] = cur_val.to_json()
-                elif isinstance(cur_val, list):
-                    json_obj[att] = []
-                    for obj in cur_val:
-                        if issubclass(type(obj), JumpStartDataHolderType):
-                            json_obj[att].append(obj.to_json())
-                        else:
-                            json_obj[att].append(obj)
-                else:
-                    json_obj[att] = cur_val
-        return json_obj

src/sagemaker/jumpstart/types.py

@@ -99,6 +99,25 @@ def __repr__(self) -> str:
         }
         return f"{type(self).__name__} at {hex(id(self))}: {str(att_dict)}"
 
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of object."""
+        json_obj = {}
+        for att in self.__slots__:
+            if hasattr(self, att):
+                cur_val = getattr(self, att)
+                if issubclass(type(cur_val), JumpStartDataHolderType):
+                    json_obj[att] = cur_val.to_json()
+                elif isinstance(cur_val, list):
+                    json_obj[att] = []
+                    for obj in cur_val:
+                        if issubclass(type(obj), JumpStartDataHolderType):
+                            json_obj[att].append(obj.to_json())
+                        else:
+                            json_obj[att].append(obj)
+                else:
+                    json_obj[att] = cur_val
+        return json_obj
+
 
 class JumpStartS3FileType(str, Enum):
     """Type of files published in JumpStart S3 distribution buckets."""
@@ -911,25 +930,6 @@ def from_hub_content_doc(self, hub_content_doc: Dict[str, Any]) -> None:
         """
         # TODO: Implement
 
-    def to_json(self) -> Dict[str, Any]:
-        """Returns json representation of JumpStartModelSpecs object."""
-        json_obj = {}
-        for att in self.__slots__:
-            if hasattr(self, att):
-                cur_val = getattr(self, att)
-                if issubclass(type(cur_val), JumpStartDataHolderType):
-                    json_obj[att] = cur_val.to_json()
-                elif isinstance(cur_val, list):
-                    json_obj[att] = []
-                    for obj in cur_val:
-                        if issubclass(type(obj), JumpStartDataHolderType):
-                            json_obj[att].append(obj.to_json())
-                        else:
-                            json_obj[att].append(obj)
-                else:
-                    json_obj[att] = cur_val
-        return json_obj
-
     def supports_prepacked_inference(self) -> bool:
         """Returns True if the model has a prepacked inference artifact."""
         return getattr(self, "hosting_prepacked_artifact_key", None) is not None