APP-6703: Added AtlanResponse and AtlanTagTranslator models

Author: Aryamanz29

pyatlan/client/asset.py

@@ -70,6 +70,7 @@
     AssetRequest,
     AssetResponse,
     AtlanObject,
+    AtlanResponse,
     AtlanTag,
     AtlanTagName,
     BulkRequest,
@@ -246,7 +247,10 @@ def search(self, criteria: IndexSearchRequest, bulk=False) -> IndexSearchResults
                     unflatten_custom_metadata_for_entity(
                         entity=entity, attributes=criteria.attributes
                     )
-                assets = parse_obj_as(List[Asset], raw_json["entities"])
+                translated_entities = AtlanResponse(
+                    raw_json=raw_json["entities"], client=self._client
+                ).to_dict()
+                assets = parse_obj_as(List[Asset], translated_entities)
             except ValidationError as err:
                 raise ErrorCode.JSON_ERROR.exception_with_parameters(
                     raw_json, 200, str(err)
@@ -318,7 +322,10 @@ def get_lineage_list(
                     unflatten_custom_metadata_for_entity(
                         entity=entity, attributes=lineage_request.attributes
                     )
-                assets = parse_obj_as(List[Asset], raw_json["entities"])
+                translated_entities = AtlanResponse(
+                    raw_json=raw_json["entities"], client=self._client
+                ).to_dict()
+                assets = parse_obj_as(List[Asset], translated_entities)
                 has_more = parse_obj_as(bool, raw_json["hasMore"])
             except ValidationError as err:
                 raise ErrorCode.JSON_ERROR.exception_with_parameters(
@@ -546,7 +553,10 @@ def _handle_relationships(self, raw_json):
                 raw_json["entity"]["relationshipAttributes"]
             )
         raw_json["entity"]["relationshipAttributes"] = {}
-        asset = AssetResponse[A](**raw_json).entity
+        translated_raw_json = AtlanResponse(
+            raw_json=raw_json, client=self._client
+        ).to_dict()
+        asset = AssetResponse[A](**translated_raw_json).entity
         asset.is_incomplete = False
         return asset
 
@@ -2082,7 +2092,10 @@ def _process_entities(self, entities):
             unflatten_custom_metadata_for_entity(
                 entity=entity, attributes=self._criteria.attributes
             )
-        self._assets = parse_obj_as(List[Asset], entities)
+        translated_entities = AtlanResponse(
+            raw_json=entities, client=self._client
+        ).to_dict()
+        self._assets = parse_obj_as(List[Asset], translated_entities)
 
     def _update_first_last_record_creation_times(self):
         self._first_record_creation_time = self._last_record_creation_time = -2

pyatlan/model/assets/core/referenceable.py

@@ -85,18 +85,21 @@ def can_be_archived(self) -> bool:
         """
         return True
 
+    # @property
+    # def atlan_tag_names(self) -> List[str]:
+    #     from pyatlan.client.atlan import AtlanClient
+    #     from pyatlan.model.constants import DELETED_
+
+    #     if self.classification_names:
+    #         return [
+    #             AtlanClient.get_current_client().atlan_tag_cache.get_name_for_id(tag_id)
+    #             or DELETED_
+    #             for tag_id in self.classification_names
+    #         ]
+    #     return []
     @property
     def atlan_tag_names(self) -> List[str]:
-        from pyatlan.client.atlan import AtlanClient
-        from pyatlan.model.constants import DELETED_
-
-        if self.classification_names:
-            return [
-                AtlanClient.get_current_client().atlan_tag_cache.get_name_for_id(tag_id)
-                or DELETED_
-                for tag_id in self.classification_names
-            ]
-        return []
+        return self.classification_names or []
 
     def __setattr__(self, name, value):
         if name in Referenceable._convenience_properties:

pyatlan/model/core.py

@@ -13,16 +13,19 @@
 
 if TYPE_CHECKING:
     from dataclasses import dataclass
+
+    from pyatlan.client.atlan import AtlanClient
 else:
     from pydantic.v1.dataclasses import dataclass
 
-from typing import Any, Dict, Generic, List, Optional, TypeVar
+from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
 
 from pydantic.v1.generics import GenericModel
 
 from pyatlan.model.constants import DELETED_, DELETED_SENTINEL
 from pyatlan.model.enums import AnnouncementType, EntityStatus, SaveSemantic
 from pyatlan.model.structs import SourceTagAttachment
+from pyatlan.model.translators import AtlanTagTranslator
 
 
 class AtlanTagName:
@@ -178,6 +181,38 @@ def from_yaml(cls, yaml_str: str):
         return cls(**data)
 
 
+class AtlanResponse:
+    def __init__(self, raw_json: Dict[str, Any], client: AtlanClient):
+        self.raw_json = raw_json
+        self.client = client
+        self.translators = [
+            AtlanTagTranslator(client),
+            # Register more translators here
+        ]
+        self.translated = self._deep_translate(self.raw_json)
+
+    def _deep_translate(
+        self, data: Union[Dict[str, Any], List[Any], Any]
+    ) -> Union[Dict[str, Any], List[Any], Any]:
+        if isinstance(data, dict):
+            # Apply translators to this dict if any apply
+            for translator in self.translators:
+                if translator.applies_to(data):
+                    data = translator.translate(data)
+
+            # Recursively apply to each value
+            return {key: self._deep_translate(value) for key, value in data.items()}
+
+        elif isinstance(data, list):
+            return [self._deep_translate(item) for item in data]
+
+        else:
+            return data
+
+    def to_dict(self) -> Dict[str, Any]:
+        return self.translated
+
+
 class SearchRequest(AtlanObject, ABC):
     attributes: Optional[List[str]] = Field(
         default_factory=list,
@@ -202,7 +237,7 @@ class AtlanTag(AtlanObject):
     class Config:
         extra = "forbid"
 
-    type_name: Optional[AtlanTagName] = Field(
+    type_name: Optional[str] = Field(
         default=None,
         description="Name of the type definition that defines this instance.\n",
         alias="typeName",
@@ -255,26 +290,26 @@ class Config:
     def source_tag_attachements(self) -> List[SourceTagAttachment]:
         return self._source_tag_attachements
 
-    @validator("type_name", pre=True)
-    def type_name_is_tag_name(cls, value):
-        if isinstance(value, AtlanTagName):
-            return value
-        return AtlanTagName._convert_to_display_text(value)
-
-    def __init__(self, *args, **kwargs):
-        from pyatlan.client.atlan import AtlanClient
-
-        super().__init__(*args, **kwargs)
-        if self.type_name != AtlanTagName.get_deleted_sentinel():
-            attr_id = AtlanClient.get_current_client().atlan_tag_cache.get_source_tags_attr_id(
-                self.type_name.id
-            )
-        if self.attributes and attr_id in self.attributes:
-            self._source_tag_attachements = [
-                SourceTagAttachment(**source_tag["attributes"])
-                for source_tag in self.attributes[attr_id]
-                if isinstance(source_tag, dict) and source_tag.get("attributes")
-            ]
+    # @validator("type_name", pre=True)
+    # def type_name_is_tag_name(cls, value):
+    #     if isinstance(value, AtlanTagName):
+    #         return value
+    #     return AtlanTagName._convert_to_display_text(value)
+
+    # def __init__(self, *args, **kwargs):
+    #     from pyatlan.client.atlan import AtlanClient
+
+    #     super().__init__(*args, **kwargs)
+    #     if self.type_name != AtlanTagName.get_deleted_sentinel():
+    #         attr_id = AtlanClient.get_current_client().atlan_tag_cache.get_source_tags_attr_id(
+    #             self.type_name.id
+    #         )
+    #     if self.attributes and attr_id in self.attributes:
+    #         self._source_tag_attachements = [
+    #             SourceTagAttachment(**source_tag["attributes"])
+    #             for source_tag in self.attributes[attr_id]
+    #             if isinstance(source_tag, dict) and source_tag.get("attributes")
+    #         ]
 
     @classmethod
     def of(

pyatlan/model/translators.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict
+
+if TYPE_CHECKING:
+    from pyatlan.client.atlan import AtlanClient
+
+
+class BaseTranslator(ABC):
+    @abstractmethod
+    def applies_to(self, data: Dict[str, Any]) -> bool:
+        pass
+
+    @abstractmethod
+    def translate(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        pass
+
+
+class AtlanTagTranslator(BaseTranslator):
+    _TYPE_NAME = "typeName"
+    _CLASSIFICATIONS = "classifications"
+    _CLASSIFICATION_NAMES = "classificationNames"
+
+    def __init__(self, client: AtlanClient):
+        self.client = client
+
+    def applies_to(self, data: Dict[str, Any]) -> bool:
+        return self._CLASSIFICATION_NAMES in data or self._CLASSIFICATIONS in data
+
+    def translate(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        raw_json = data.copy()
+
+        if self._CLASSIFICATION_NAMES in raw_json:
+            raw_json[self._CLASSIFICATION_NAMES] = [
+                self.client.atlan_tag_cache.get_name_for_id(tag_id)
+                for tag_id in raw_json[self._CLASSIFICATION_NAMES]
+            ]
+
+        if self._CLASSIFICATIONS in raw_json:
+            for classification in raw_json[self._CLASSIFICATIONS]:
+                tag_id = classification.get(self._TYPE_NAME)
+                if tag_id:
+                    classification[self._TYPE_NAME] = (
+                        self.client.atlan_tag_cache.get_name_for_id(tag_id)
+                    )
+
+        return raw_json