Refactored to support type hints for MutatedEntities

Author: ErnestoLoma

pyatlan/client/entity.py

@@ -17,15 +17,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Type, TypeVar
+from typing import Type, TypeVar, Union
 
 from pyatlan.client.atlan import AtlanClient
-from pyatlan.model.core import (
-    AssetResponse,
+from pyatlan.model.assets import (
+    Asset,
+    Referenceable,
+    AtlasGlossary,
+    AtlasGlossaryCategory,
+    AtlasGlossaryTerm,
     AssetMutationResponse,
-    BulkRequest,
 )
-
+from pyatlan.model.core import AssetResponse, BulkRequest
 from pyatlan.model.enums import AtlanDeleteType
 from pyatlan.utils import (
     API,
@@ -37,7 +40,7 @@
     HTTPStatus,
 )
 
-T = TypeVar("T")
+T = TypeVar("T", bound=Referenceable)
 
 
 class EntityClient:
@@ -211,13 +214,18 @@ class EntityClient:
     def __init__(self, client: AtlanClient):
         self.client = client
 
+    Assets = Union[AtlasGlossary, AtlasGlossaryCategory, AtlasGlossaryTerm]
+    Asset_Types = Union[
+        Type[AtlasGlossary], Type[AtlasGlossaryCategory], Type[AtlasGlossaryTerm]
+    ]
+
     def get_entity_by_guid(
         self,
         guid,
-        asset_type: Type[T],
+        asset_type: Asset_Types,
         min_ext_info: bool = False,
         ignore_relationships: bool = False,
-    ) -> T:
+    ) -> Assets:
         query_params = {
             "minExtInfo": min_ext_info,
             "ignoreRelationships": ignore_relationships,
@@ -231,25 +239,21 @@ def get_entity_by_guid(
             raw_json["entity"]["relationshipAttributes"]
         )
         raw_json["entity"]["relationshipAttributes"] = {}
-        response = AssetResponse[asset_type](**raw_json)
-        return response.entity
-
-    def update_entity(self, entity: T) -> AssetMutationResponse[T]:
-        request = BulkRequest[T](entities=[entity])
-        raw_json = self.client.call_api(EntityClient.BULK_UPDATE, None, request)
-        response: AssetMutationResponse[T] = AssetMutationResponse(**raw_json)
-        return response
+        if issubclass(asset_type, AtlasGlossary):
+            return AssetResponse[AtlasGlossary](**raw_json).entity
+        if issubclass(asset_type, AtlasGlossaryCategory):
+            return AssetResponse[AtlasGlossaryCategory](**raw_json).entity
+        if issubclass(asset_type, AtlasGlossaryTerm):
+            return AssetResponse[AtlasGlossaryTerm](**raw_json).entity
 
-    def create_entity(self, entity: T) -> AssetMutationResponse[T]:
+    def upsert(self, entity: Asset) -> AssetMutationResponse:
         request = BulkRequest[T](entities=[entity])
         raw_json = self.client.call_api(EntityClient.BULK_UPDATE, None, request)
-        response: AssetMutationResponse[T] = AssetMutationResponse(**raw_json)
-        return response
+        return AssetMutationResponse(**raw_json)
 
-    def purge_entity_by_guid(self, guid) -> AssetMutationResponse[T]:
+    def purge_entity_by_guid(self, guid) -> AssetMutationResponse:
         raw_json = self.client.call_api(
             EntityClient.DELETE_ENTITY_BY_GUID.format_path_with_params(guid),
             {"deleteType": AtlanDeleteType.HARD.value},
         )
-        response: AssetMutationResponse[T] = AssetMutationResponse(**raw_json)
-        return response
+        return AssetMutationResponse(**raw_json)

pyatlan/generator/templates/entity.jinja2

@@ -24,7 +24,7 @@ class {{struct.name}}(AtlanObject):
 {% endfor %}
 {% for entity_def in entity_defs %}
 {%- set super_classes = ['AtlanObject'] if not entity_def.super_types else entity_def.super_types -%}
-class {{ entity_def.name }}({{super_classes[0]}}):
+class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes  %}, type_name='{{ entity_def.name }}'{% endif %}):
     """Description"""
 {%- if entity_def.name == "Referenceable"  %}
     class Attributes(AtlanObject):
@@ -57,6 +57,11 @@ class {{ entity_def.name }}({{super_classes[0]}}):
         description='Time (epoch) at which this object was created, in milliseconds.\n',
         example=1648852296555,
     )
+    delete_handler: Optional[str] = Field(
+        None,
+        description="Details on the handler used for deletion of the asset.",
+        example="Hard",
+    )
     guid: Optional[str] = Field(
         description='Unique identifier of the entity instance.\n',
         example='917ffec9-fa84-4c59-8e6c-c7b114d04be3',
@@ -123,6 +128,36 @@ class {{ entity_def.name }}({{super_classes[0]}}):
         None, description="", alias="meanings"
     )
 {%- else %}
+    {%- if entity_def.name == "Asset"  %}
+    _subtypes_:dict[str, type] = dict()
+
+    def __init_subclass__(cls, type_name=None):
+        cls._subtypes_[type_name or cls.__name__.lower()] = cls
+
+    @classmethod
+    def __get_validators__(cls):
+        yield cls._convert_to_real_type_
+
+    @classmethod
+    def _convert_to_real_type_(cls, data):
+
+        if isinstance(data, Asset):
+            return data
+
+        data_type = (
+            data.get("type_name") if "type_name" in data else data.get("typeName")
+        )
+
+        if data_type is None:
+            raise ValueError("Missing 'type' in Asset")
+
+        sub = cls._subtypes_.get(data_type)
+
+        if sub is None:
+            raise TypeError(f"Unsupport sub-type: {data_type}")
+
+        return sub(**data)
+    {%- endif %}
     {%- if entity_def.attribute_defs %}
         {%- if not entity_def.sub_types %}
 
@@ -176,6 +211,34 @@ class {{ entity_def.name }}({{super_classes[0]}}):
     {% endif %}
 {%- endif %}
 {% endfor %}
+class MutatedEntities(AtlanObject):
+    CREATE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were created. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="CREATE",
+    )
+    UPDATE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were updated. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="UPDATE",
+    )
+    DELETE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were deleted. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="DELETE",
+    )
+
+
+class AssetMutationResponse(AtlanObject):
+    guid_assignments: dict[str, Any] = Field(
+        None, description="Map of assigned unique identifiers for the changed assets."
+    )
+    mutated_entities: Optional[MutatedEntities] = Field(
+        None, description="Assets that were changed."
+    )
 Referenceable.update_forward_refs()
 AtlasGlossary.update_forward_refs()
 {% for entity_def in entity_defs %}

pyatlan/model/assets.py

@@ -177,6 +177,11 @@ class Attributes(AtlanObject):
         description="Time (epoch) at which this object was created, in milliseconds.\n",
         example=1648852296555,
     )
+    delete_handler: Optional[str] = Field(
+        None,
+        description="Details on the handler used for deletion of the asset.",
+        example="Hard",
+    )
     guid: Optional[str] = Field(
         description="Unique identifier of the entity instance.\n",
         example="917ffec9-fa84-4c59-8e6c-c7b114d04be3",
@@ -245,6 +250,35 @@ class Attributes(AtlanObject):
 class Asset(Referenceable):
     """Description"""
 
+    _subtypes_: dict[str, type] = dict()
+
+    def __init_subclass__(cls, type_name=None):
+        cls._subtypes_[type_name or cls.__name__.lower()] = cls
+
+    @classmethod
+    def __get_validators__(cls):
+        yield cls._convert_to_real_type_
+
+    @classmethod
+    def _convert_to_real_type_(cls, data):
+
+        if isinstance(data, Asset):
+            return data
+
+        data_type = (
+            data.get("type_name") if "type_name" in data else data.get("typeName")
+        )
+
+        if data_type is None:
+            raise ValueError("Missing 'type' in Asset")
+
+        sub = cls._subtypes_.get(data_type)
+
+        if sub is None:
+            raise TypeError(f"Unsupport sub-type: {data_type}")
+
+        return sub(**data)
+
     class Attributes(Referenceable.Attributes):
         name: str = Field(None, description="", alias="name")
         display_name: Optional[str] = Field(None, description="", alias="displayName")
@@ -558,7 +592,7 @@ def clear_announcment(self):
         self.attributes.announcement_type = None
 
 
-class AtlasGlossary(Asset):
+class AtlasGlossary(Asset, type_name="AtlasGlossary"):
     """Description"""
 
     type_name: Literal["AtlasGlossary"] = Field("AtlasGlossary")
@@ -601,15 +635,15 @@ class Attributes(Asset.Attributes):
     )
 
 
-class DataSet(Asset):
+class DataSet(Asset, type_name="DataSet"):
     """Description"""
 
 
-class ProcessExecution(Asset):
+class ProcessExecution(Asset, type_name="ProcessExecution"):
     """Description"""
 
 
-class AtlasGlossaryTerm(Asset):
+class AtlasGlossaryTerm(Asset, type_name="AtlasGlossaryTerm"):
     """Description"""
 
     type_name: Literal["AtlasGlossaryTerm"] = Field("AtlasGlossaryTerm")
@@ -695,15 +729,15 @@ class Attributes(Asset.Attributes):
     )
 
 
-class Cloud(Asset):
+class Cloud(Asset, type_name="Cloud"):
     """Description"""
 
 
-class Infrastructure(Asset):
+class Infrastructure(Asset, type_name="Infrastructure"):
     """Description"""
 
 
-class Connection(Asset):
+class Connection(Asset, type_name="Connection"):
     """Description"""
 
     type_name: Literal["Connection"] = Field("Connection")
@@ -770,7 +804,7 @@ class Attributes(Asset.Attributes):
     )
 
 
-class Process(Asset):
+class Process(Asset, type_name="Process"):
     """Description"""
 
     class Attributes(Asset.Attributes):
@@ -802,7 +836,7 @@ class Attributes(Asset.Attributes):
     )
 
 
-class AtlasGlossaryCategory(Asset):
+class AtlasGlossaryCategory(Asset, type_name="AtlasGlossaryCategory"):
     """Description"""
 
     type_name: Literal["AtlasGlossaryCategory"] = Field("AtlasGlossaryCategory")
@@ -849,7 +883,7 @@ class Attributes(Asset.Attributes):
     )
 
 
-class Badge(Asset):
+class Badge(Asset, type_name="Badge"):
     """Description"""
 
     type_name: Literal["Badge"] = Field("Badge")
@@ -881,11 +915,11 @@ class Attributes(Asset.Attributes):
     )
 
 
-class Namespace(Asset):
+class Namespace(Asset, type_name="Namespace"):
     """Description"""
 
 
-class Catalog(Asset):
+class Catalog(Asset, type_name="Catalog"):
     """Description"""
 
 
@@ -5566,6 +5600,36 @@ class Attributes(Salesforce.Attributes):
     )
 
 
+class MutatedEntities(AtlanObject):
+    CREATE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were created. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="CREATE",
+    )
+    UPDATE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were updated. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="UPDATE",
+    )
+    DELETE: Optional[list[Asset]] = Field(
+        None,
+        description="Assets that were deleted. The detailed properties of the returned asset will vary based on the "
+        "type of asset, but listed in the example are the common set of properties across assets.",
+        alias="DELETE",
+    )
+
+
+class AssetMutationResponse(AtlanObject):
+    guid_assignments: dict[str, Any] = Field(
+        None, description="Map of assigned unique identifiers for the changed assets."
+    )
+    mutated_entities: Optional[MutatedEntities] = Field(
+        None, description="Assets that were changed."
+    )
+
+
 Referenceable.update_forward_refs()
 AtlasGlossary.update_forward_refs()
 

pyatlan/model/core.py

@@ -101,33 +101,3 @@ class AssetRequest(AtlanObject, GenericModel, Generic[T]):
 
 class BulkRequest(AtlanObject, GenericModel, Generic[T]):
     entities: list[T]
-
-
-class MutatedEntities(AtlanObject, GenericModel, Generic[T]):
-    CREATE: Optional[list[T]] = Field(
-        None,
-        description="Assets that were created. The detailed properties of the returned asset will vary based on the "
-        "type of asset, but listed in the example are the common set of properties across assets.",
-        alias="CREATE",
-    )
-    UPDATE: Optional[list[T]] = Field(
-        None,
-        description="Assets that were updated. The detailed properties of the returned asset will vary based on the "
-        "type of asset, but listed in the example are the common set of properties across assets.",
-        alias="UPDATE",
-    )
-    DELETE: Optional[list[T]] = Field(
-        None,
-        description="Assets that were deleted. The detailed properties of the returned asset will vary based on the "
-        "type of asset, but listed in the example are the common set of properties across assets.",
-        alias="DELETE",
-    )
-
-
-class AssetMutationResponse(AtlanObject, GenericModel, Generic[T]):
-    guid_assignments: dict[str, Any] = Field(
-        None, description="Map of assigned unique identifiers for the changed assets."
-    )
-    mutated_entities: Optional[MutatedEntities] = Field(
-        None, description="Assets that were changed."
-    )