[fix] Fixed source tag translations

Author: Aryamanz29

pyatlan/client/asset.py

@@ -640,7 +640,10 @@ def save(
             asset.validate_required()
         request = BulkRequest[Asset](entities=entities)
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
-        response = AssetMutationResponse(**raw_json)
+        translated_raw_json = AtlanResponse(
+            raw_json=raw_json, client=self._client
+        ).to_dict()
+        response = AssetMutationResponse(**translated_raw_json)
         if connections_created := response.assets_created(Connection):
             self._wait_for_connections_to_be_created(connections_created)
         return response

pyatlan/model/core.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING
 
 import yaml  # type: ignore[import-untyped]
-from pydantic.v1 import BaseModel, Extra, Field, PrivateAttr, root_validator, validator
+from pydantic.v1 import BaseModel, Extra, Field, root_validator, validator
 
 from pyatlan.model.utils import encoders, to_camel_case
 
@@ -280,15 +280,15 @@ class Config:
         alias="restrictPropagationThroughHierarchy",
     )
     validity_periods: Optional[List[str]] = Field(default=None, alias="validityPeriods")
-    _source_tag_attachements: List[SourceTagAttachment] = PrivateAttr(
-        default_factory=list
+    source_tag_attachements: List[SourceTagAttachment] = Field(
+        default_factory=list, exclude=True
     )
 
     attributes: Optional[Dict[str, Any]] = None
 
-    @property
-    def source_tag_attachements(self) -> List[SourceTagAttachment]:
-        return self._source_tag_attachements
+    # @property
+    # 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):
@@ -340,7 +340,7 @@ def of(
                 or ""
             )
             tag.attributes = {source_tag_attr_id: [source_tag_attachment]}
-            tag._source_tag_attachements.append(source_tag_attachment)
+            tag.source_tag_attachements.append(source_tag_attachment)
         return tag
 
 

pyatlan/model/translators.py

@@ -3,6 +3,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict
 
+from pyatlan.model.structs import SourceTagAttachment
+
 if TYPE_CHECKING:
     from pyatlan.client.atlan import AtlanClient
 
@@ -41,8 +43,22 @@ def translate(self, data: Dict[str, Any]) -> Dict[str, Any]:
             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)
+                    tag_name = self.client.atlan_tag_cache.get_name_for_id(tag_id)
+                    if not tag_name:
+                        return
+                    classification[self._TYPE_NAME] = tag_name
+                    # Check if the tag is a source tag (in that case tag has "attributes")
+                    attr_id = self.client.atlan_tag_cache.get_source_tags_attr_id(
+                        tag_id
                     )
+                    if attr_id:
+                        classification["source_tag_attachements"] = [
+                            SourceTagAttachment(**source_tag["attributes"])
+                            for source_tag in classification.get("attributes").get(
+                                attr_id
+                            )
+                            if isinstance(source_tag, dict)
+                            and source_tag.get("attributes")
+                        ]
 
         return raw_json