[generator + tests] Updated generator scripts + added unit tests

Author: Aryamanz29

pyatlan/generator/class_generator.py

@@ -18,8 +18,8 @@
 import networkx as nx
 from jinja2 import Environment, PackageLoader
 
-from pyatlan.model.typedef import EntityDef, EnumDef, TypeDefResponse
-from pyatlan.model.utils import to_snake_case
+from pyatlan.model.typedef import EntityDef, EnumDef, RelationshipDef, TypeDefResponse
+from pyatlan.model.utils import to_python_class_name, to_snake_case
 
 REFERENCEABLE = "Referenceable"
 TYPE_DEF_FILE = Path(os.getenv("TMPDIR", "/tmp")) / "typedefs.json"
@@ -80,6 +80,7 @@
 # across Python versions (e.g: 3.8 and 3.9).
 PARENT = Path(__file__).resolve().parent
 ASSETS_DIR = PARENT.parent / "model" / "assets"
+ASSETS_RELATIONS_DIR = PARENT.parent / "model" / "assets" / "relations"
 CORE_ASSETS_DIR = PARENT.parent / "model" / "assets" / "core"
 MODEL_DIR = PARENT.parent / "model"
 DOCS_DIR = PARENT.parent / "documentation"
@@ -450,6 +451,29 @@ def create_modules(cls):
                             cls._CORE_ASSETS.add(related_asset.super_class)
 
 
+class RelationshipDefInfo:
+    module_name: str
+    relationship_def_infos: List["RelationshipDefInfo"] = []
+
+    def __init__(self, name: str, relationship_def: RelationshipDef):
+        self.module_name = to_snake_case(name)
+        self.relationship_def = relationship_def
+
+    @classmethod
+    def create(cls, relationship_defs: List[RelationshipDef]):
+        for rel_def in relationship_defs:
+            # Only pick `atlas_core` enums, not user-created ones.
+            if rel_def.attribute_defs:
+                cls.relationship_def_infos.append(
+                    RelationshipDefInfo(
+                        name=to_snake_case(rel_def.name), relationship_def=rel_def
+                    )
+                )
+                # if rel_def.name == "UserDefRelationship":
+                #     import ipdb; ipdb.set_trace()
+                # print(to_python_class_name(rel_def.name))
+
+
 class AttributeType(Enum):
     PRIMITIVE = "PRIMITIVE"
     ENUM = "ENUM"
@@ -664,6 +688,7 @@ def __init__(self) -> None:
         )
         self.environment.filters["to_snake_case"] = to_snake_case
         self.environment.filters["get_type"] = get_type
+        self.environment.filters["to_cls_name"] = to_python_class_name
         self.environment.filters["get_search_type"] = get_search_type
         self.environment.filters["get_mapped_type"] = get_mapped_type
         self.environment.filters["get_class_var_for_attr"] = get_class_var_for_attr
@@ -719,6 +744,21 @@ def render_module(self, asset_info: AssetInfo, enum_defs: List["EnumDefInfo"]):
         with (ASSETS_DIR / f"{asset_info.module_name}.py").open("w") as script:
             script.write(content)
 
+    def render_custom_relationship_module(
+        self, relationship_def_info: RelationshipDefInfo
+    ):
+        template = self.environment.get_template("custom_relationship.jinja2")
+        content = template.render(
+            {
+                "relationship_info": relationship_info.relationship_def,
+                "templates_path": TEMPLATES_DIR.absolute().as_posix(),
+            }
+        )
+        with (ASSETS_RELATIONS_DIR / f"{relationship_def_info.module_name}.py").open(
+            "w"
+        ) as script:
+            script.write(content)
+
     def render_core_module(self, asset_info: AssetInfo, enum_defs: List["EnumDefInfo"]):
         template = self.environment.get_template("module.jinja2")
         content = template.render(
@@ -769,6 +809,22 @@ def render_mypy_init(self, assets: List[AssetInfo]):
         with init_path.open("w") as script:
             script.write(content)
 
+    def render_relations_init(self, relation_def_infos: List[RelationshipDefInfo]):
+        template = self.environment.get_template("relations_init.jinja2")
+        content = template.render({"relation_def_infos": relation_def_infos})
+
+        init_path = ASSETS_RELATIONS_DIR / "__init__.py"
+        with init_path.open("w") as script:
+            script.write(content)
+
+    def render_relations_mypy_init(self, relation_def_infos: List[RelationshipDefInfo]):
+        template = self.environment.get_template("relations_mypy_init.jinja2")
+        content = template.render({"relation_def_infos": relation_def_infos})
+
+        init_path = ASSETS_RELATIONS_DIR / "__init__.pyi"
+        with init_path.open("w") as script:
+            script.write(content)
+
     def render_structs(self, struct_defs):
         template = self.environment.get_template("structs.jinja2")
         content = template.render({"struct_defs": struct_defs})
@@ -933,8 +989,10 @@ def filter_attributes_of_custom_entity_type():
     filter_attributes_of_custom_entity_type()
     AssetInfo.sub_type_names_to_ignore = type_defs.custom_entity_def_names
     AssetInfo.set_entity_defs(type_defs.reserved_entity_defs)
+    RelationshipDefInfo.create(type_defs.relationship_defs)
     AssetInfo.update_all_circular_dependencies()
     AssetInfo.create_modules()
+
     for file in (ASSETS_DIR).glob("*.py"):
         file.unlink()
     for file in (CORE_ASSETS_DIR).glob("*.py"):
@@ -944,14 +1002,28 @@ def filter_attributes_of_custom_entity_type():
         file.unlink()
     generator = Generator()
     EnumDefInfo.create(type_defs.enum_defs)
+
     for asset_info in ModuleInfo.assets.values():
         if asset_info.is_core_asset or asset_info.name in asset_info._CORE_ASSETS:
             generator.render_core_module(asset_info, EnumDefInfo.enum_def_info)
         else:
             generator.render_module(asset_info, EnumDefInfo.enum_def_info)
+
+    for file in (ASSETS_RELATIONS_DIR).glob("*.py"):
+        # Remove all files in the core assets
+        # directory except `relationship_attributes.py`
+        if file.name == "relationship_attributes.py":
+            continue
+        file.unlink()
+
+    for relationship_info in RelationshipDefInfo.relationship_def_infos:
+        generator.render_custom_relationship_module(relationship_info)
+
     generator.render_init(ModuleInfo.assets.values())  # type: ignore
     generator.render_core_init(ModuleInfo.assets.values())  # type: ignore
     generator.render_mypy_init(ModuleInfo.assets.values())  # type: ignore
+    generator.render_relations_init(RelationshipDefInfo.relationship_def_infos)  # type: ignore
+    generator.render_relations_mypy_init(RelationshipDefInfo.relationship_def_infos)  # type: ignore
     generator.render_structs(type_defs.struct_defs)
     generator.render_enums(EnumDefInfo.enum_def_info)
     generator.render_docs_struct_snippets(type_defs.struct_defs)

pyatlan/generator/templates/custom_relationship.jinja2

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from pydantic.v1 import Field, validator
+
+from pyatlan.model.assets import Asset
+from pyatlan.model.assets.relations import RelationshipAttributes
+from pyatlan.model.core import AtlanObject
+from pyatlan.model.enums import SaveSemantic, AtlasGlossaryTermRelationshipStatus
+
+
+class {{ relationship_info.name | to_cls_name }}(RelationshipAttributes):
+    type_name: str = Field(
+        allow_mutation=False,
+        default="{{ relationship_info.name }}",
+        description="{{ relationship_info.description }}",
+    )
+    attributes: {{ relationship_info.name | to_cls_name }}.Attributes = Field(
+        default_factory=lambda: {{ relationship_info.name | to_cls_name }}.Attributes(),
+        description="Map of attributes in the instance and their values",
+    )
+
+    class Attributes(AtlanObject):
+        {% for attribute_def in relationship_info.attribute_defs -%}
+        {{ attribute_def["name"] | to_snake_case }}: Optional[{{ attribute_def["typeName"] | get_type }}] = Field(
+            default=None,
+            description="{{ attribute_def["description"] }}",
+        )
+        {% endfor %}
+
+    def __init__(__pydantic_self__, **data: Any) -> None:
+        if "attributes" not in data:
+            data = {"attributes": data}
+        super().__init__(**data)
+        __pydantic_self__.__fields_set__.update(["attributes", "type_name"])
+
+    class {{ relationship_info.end_def1["name"] | to_cls_name }}(Asset):
+        type_name: str = Field(
+            default="{{ relationship_info.name }}",
+            description="{{ relationship_info.end_def1["description"] or 'Name of the relationship type that defines the relationship.'}}",
+        )
+        relationship_type: str = Field(
+            default="{{ relationship_info.name }}",
+            description="Fixed typeName for {{ relationship_info.name }}.",
+        )
+        relationship_attributes: {{ relationship_info.name | to_cls_name }} = Field(
+            default=None,
+            description="Attributes of the {{ relationship_info.name }}.",
+        )
+
+        @validator("type_name")
+        def validate_type_name(cls, v):
+            return v
+
+        def __init__(__pydantic_self__, **data: Any) -> None:
+            super().__init__(**data)
+            __pydantic_self__.__fields_set__.update(["type_name", "relationship_type"])
+
+    class {{ relationship_info.end_def2["name"] | to_cls_name }}(Asset):
+        type_name: str = Field(
+            default="{{ relationship_info.name }}",
+            description="{{ relationship_info.end_def2["description"] or 'Name of the relationship type that defines the relationship.'}}",
+        )
+        relationship_type: str = Field(
+            default="{{ relationship_info.name }}",
+            description="Fixed typeName for {{ relationship_info.name }}.",
+        )
+        relationship_attributes: {{ relationship_info.name | to_cls_name }} = Field(
+            default=None,
+            description="Attributes of the {{ relationship_info.name }}.",
+        )
+
+        @validator("type_name")
+        def validate_type_name(cls, v):
+            return v
+
+        def __init__(__pydantic_self__, **data: Any) -> None:
+            super().__init__(**data)
+            __pydantic_self__.__fields_set__.update(["type_name", "relationship_type"])
+
+    def {{ relationship_info.end_def1["name"] | to_snake_case }}(
+        self, related: Asset, semantic: SaveSemantic = SaveSemantic.REPLACE
+    ) -> {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def1["name"] | to_cls_name }}:
+        if related.guid:
+            return {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def1["name"] | to_cls_name }}._create_ref(
+                type_name=related.type_name,
+                guid=related.guid,
+                semantic=semantic,
+                relationship_attributes=self,
+            )
+
+        # If the related asset does not have a GUID, we use qualifiedName
+        return {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def1["name"] | to_cls_name }}._create_ref(
+            type_name=related.type_name,
+            unique_attributes={"qualifiedName": related.qualified_name},
+            semantic=semantic,
+            relationship_attributes=self,
+        )
+
+    def {{ relationship_info.end_def2["name"] | to_snake_case }}(
+        self, related: Asset, semantic: SaveSemantic = SaveSemantic.REPLACE
+    ) -> {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def2["name"] | to_cls_name }}:
+        if related.guid:
+            return {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def2["name"] | to_cls_name }}._create_ref(
+                type_name=related.type_name,
+                guid=related.guid,
+                semantic=semantic,
+                relationship_attributes=self,
+            )
+
+        # If the related asset does not have a GUID, we use qualifiedName
+        return {{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def2["name"] | to_cls_name }}._create_ref(
+            type_name=related.type_name,
+            unique_attributes={"qualifiedName": related.qualified_name},
+            semantic=semantic,
+            relationship_attributes=self,
+        )
+
+
+{{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def1["name"] | to_cls_name }}.update_forward_refs()
+{{ relationship_info.name | to_cls_name }}.{{ relationship_info.end_def2["name"] | to_cls_name }}.update_forward_refs()
+{{ relationship_info.name | to_cls_name }}.update_forward_refs()

pyatlan/generator/templates/imports.jinja2

@@ -127,3 +127,4 @@ from pyatlan.model.data_mesh import DataProductsAssetsDSL
 from pyatlan.model.contract import DataContractSpec
 from pyatlan.model.lineage_ref import LineageRef
 from pyatlan.model.utils import construct_object_key
+from pyatlan.model.assets.relations import RelationshipAttributes

pyatlan/generator/templates/referenceable_attributes.jinja2

@@ -9,6 +9,11 @@
         {%- set type = attribute_def.typeName | get_type %}
         {{attribute_def.name | to_snake_case }}: {% if attribute_def.isOptional %}Optional[{% endif %}{{type}}{% if attribute_def.isOptional %}]{% endif %} = Field({% if attribute_def.isOptional %}default=None,{% endif %} description='') # relationship
         {%- endfor %}
+        relationship_attributes: Optional[RelationshipAttributes] = Field(
+            default=None,
+            description="Map of relationships for the entity. The specific keys of this map will vary by type, "
+            "so are described in the sub-types of this schema.",
+        )
 
         def validate_required(self):
             pass
@@ -122,10 +127,10 @@
     )
     is_incomplete: Optional[bool] = Field(default=None, description="", example=True)
     labels: Optional[List[str]] = Field(default=None, description='Arbitrary textual labels for the asset.')
-    relationship_attributes: Optional[Dict[str, Any]] = Field(
+    relationship_attributes: Optional[RelationshipAttributes] = Field(
         default=None,
-        description='Map of relationships for the entity. The specific keys of this map will vary by type, '
-                    'so are described in the sub-types of this schema.',
+        description="Map of relationships for the entity. The specific keys of this map will vary by type, "
+        "so are described in the sub-types of this schema.",
     )
     status: Optional[EntityStatus] = Field(
         default=None,

pyatlan/generator/templates/relations_init.jinja2

@@ -0,0 +1,13 @@
+# Copyright 2025 Atlan Pte. Ltd.
+# isort: skip_file
+import lazy_loader as lazy
+
+__PYATLAN_ASSET_RELATIONS__ = {
+    "relationship_attributes": ["RelationshipAttributes"],
+    {% for relation in relation_def_infos -%}
+    "{{ relation.relationship_def.name | to_snake_case }}": ["{{ relation.relationship_def.name | to_cls_name}}"],
+    {% endfor %}
+}
+
+lazy_loader = lazy.attach(__name__, submod_attrs=__PYATLAN_ASSET_RELATIONS__)
+__getattr__, __dir__, __all__ = lazy_loader

pyatlan/generator/templates/relations_mypy_init.jinja2

@@ -0,0 +1,13 @@
+# Copyright 2025 Atlan Pte. Ltd.
+
+__all__ = [
+    "RelationshipAttributes",
+    {% for relation in relation_def_infos -%}
+    "{{ relation.relationship_def.name | to_cls_name}}",
+    {% endfor %}
+]
+
+from .relationship_attributes import RelationshipAttributes
+{% for relation in relation_def_infos -%}
+from .{{ relation.relationship_def.name | to_snake_case }} import {{ relation.relationship_def.name | to_cls_name }}
+{% endfor %}

pyatlan/generator/templates/structs.jinja2

@@ -13,7 +13,6 @@ from pyatlan.model.enums import (
     BadgeComparisonOperator,
     BadgeConditionColor,
     SourceCostUnitType,
-    alpha_DQRuleThresholdUnit,
     FormFieldDimension,
     FormFieldType
 )