Merge pull request #415 from atlanhq/DVX-761

DVX-761: Fixed generator to handle new core assets in the typedefs

docs/asset/custom.rst

@@ -0,0 +1,10 @@
+.. _custom:
+
+Custom
+======
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: Custom
+    :members:

docs/asset/customdataset.rst

@@ -0,0 +1,10 @@
+.. _customdataset:
+
+CustomDataset
+=============
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: CustomDataset
+    :members:

docs/asset/customfield.rst

@@ -0,0 +1,10 @@
+.. _customfield:
+
+CustomField
+===========
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: CustomField
+    :members:

docs/asset/customtable.rst

@@ -0,0 +1,10 @@
+.. _customtable:
+
+CustomTable
+===========
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: CustomTable
+    :members:

docs/asset/detectidata.rst

@@ -0,0 +1,10 @@
+.. _detectidata:
+
+DetectiData
+===========
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DetectiData
+    :members:

docs/asset/detectidatadossier.rst

@@ -0,0 +1,10 @@
+.. _detectidatadossier:
+
+DetectiDataDossier
+==================
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DetectiDataDossier
+    :members:

docs/asset/detectidatadossierelement.rst

@@ -0,0 +1,10 @@
+.. _detectidatadossierelement:
+
+DetectiDataDossierElement
+=========================
+
+.. module:: pyatlan.model.assets
+    :no-index:
+
+.. autoclass:: DetectiDataDossierElement
+    :members:

docs/assets.rst

@@ -81,6 +81,10 @@ You can interact with all of the following different kinds of assets:
     asset/cubedimension
     asset/cubefield
     asset/cubehierarchy
+    asset/custom
+    asset/customdataset
+    asset/customfield
+    asset/customtable
     asset/datacontract
     asset/datadomain
     asset/datamesh

pyatlan/generator/class_generator.py

@@ -224,6 +224,7 @@ class AssetInfo:
         "MatillionGroup",
         "Stakeholder",
         "StakeholderTitle",
+        "NoSQL",
     }
 
     def __init__(self, name: str, entity_def: EntityDef):
@@ -259,7 +260,13 @@ def import_super_class(self):
         if self._name == REFERENCEABLE:
             return ""
         super_type = AssetInfo.asset_info_by_name[self.entity_def.super_types[0]]
-        if not self.is_core_asset and super_type.is_core_asset:
+        if self.name not in self._CORE_ASSETS and super_type.name in self._CORE_ASSETS:
+            return f"from .core.{super_type.module_name} import {super_type.name}"
+        elif (
+            not self.is_core_asset
+            and super_type.is_core_asset
+            and self.name not in self._CORE_ASSETS
+        ):
             return f"from .core.{super_type.module_name} import {super_type.name}"
         else:
             return f"from .{super_type.module_name} import {super_type.name}"
@@ -725,11 +732,15 @@ def render_init(self, assets: List[AssetInfo]):
             script.write(content)
 
     def render_core_init(self, assets: List[AssetInfo]):
-        asset_names = [asset.name for asset in assets if asset.is_core_asset]
+        asset_names = [
+            asset.name
+            for asset in assets
+            if asset.is_core_asset or asset.name in asset._CORE_ASSETS
+        ]
         asset_imports = [
             f"from .{asset.module_name} import {asset.name}"
             for asset in assets
-            if asset.is_core_asset
+            if asset.is_core_asset or asset.name in asset._CORE_ASSETS
         ]
 
         template = self.environment.get_template("core/init.jinja2")
@@ -922,7 +933,7 @@ def filter_attributes_of_custom_entity_type():
     generator = Generator()
     EnumDefInfo.create(type_defs.enum_defs)
     for asset_info in ModuleInfo.assets.values():
-        if asset_info.is_core_asset:
+        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)

pyatlan/generator/templates/imports.jinja2

@@ -104,6 +104,7 @@ from pyatlan.model.structs import (
     SourceTagAttribute,
     StarredDetails,
     DbtJobRun,
+    CustomRatings,
 )
 from pyatlan.utils import (
     init_guid,

pyatlan/generator/templates/init.jinja2

@@ -5,13 +5,13 @@ import lazy_loader as lazy
 __PYATLAN_ASSETS__ = {
     "core": [
         {% for asset in assets -%}
-            {% if asset.is_core_asset %}
+            {% if asset.is_core_asset or asset.name in asset._CORE_ASSETS %}
                 "{{ asset.name }}"{% if not loop.last %},{% endif %}
             {% endif %}
         {% endfor %}
     ],
     {% for asset in assets -%}
-        {% if not asset.is_core_asset %}
+        {% if not asset.is_core_asset and asset.name not in asset._CORE_ASSETS  %}
         "{{ asset.module_name}}": ["{{ asset.name }}"{% if not loop.last %}{% endif %}],
         {% endif %}
     {% endfor %}

pyatlan/generator/templates/module.jinja2

@@ -97,6 +97,6 @@ from .s_q_l import SQL  # noqa # isort:skip
 {{ import }}
 {% endfor %}
 
-{%  if not asset_info.is_core_asset %}
+{% if asset_info.name not in asset_info._CORE_ASSETS %}
 {{ asset_info.name }}.Attributes.update_forward_refs()
 {% endif %}

pyatlan/generator/templates/mypy_init.jinja2

@@ -2,7 +2,7 @@
 
 __all__ = [
     {% for asset in assets -%}
-        {% if asset.is_core_asset %}
+        {% if asset.is_core_asset or asset.name in asset._CORE_ASSETS %}
             "{{ asset.name }}"{% if not loop.last %},{% endif %}
         {% endif %}
     {% endfor %}
@@ -14,7 +14,7 @@ __all__ = [
 ]
 
 {% for asset in assets -%}
-{% if asset.is_core_asset %}
+{% if asset.is_core_asset or asset.name in asset._CORE_ASSETS %}
 from .core.{{ asset.module_name }} import {{ asset.name }}
 {% else %}
 from .{{ asset.module_name }} import {{ asset.name }}

pyatlan/model/assets/__init__.py

@@ -25,8 +25,10 @@
         "DataQuality",
         "BI",
         "Resource",
+        "Custom",
         "DataMesh",
         "SQL",
+        "NoSQL",
         "Matillion",
         "Dbt",
         "Model",
@@ -62,11 +64,11 @@
         "TablePartition",
         "Column",
         "SnowflakeStream",
-        "DatabricksUnityCatalogTag",
         "CalculationView",
         "Database",
         "Procedure",
         "SnowflakeTag",
+        "CosmosMongoDB",
         "MatillionGroup",
         "MatillionJob",
         "MatillionProject",
@@ -101,7 +103,12 @@
         "PowerBIDataflow",
         "PowerBIPage",
         "SnowflakeDynamicTable",
+        "MongoDBCollection",
         "DynamoDBSecondaryIndex",
+        "MongoDBDatabase",
+        "CosmosMongoDBAccount",
+        "CosmosMongoDBCollection",
+        "CosmosMongoDBDatabase",
     ],
     "task": ["Task"],
     "data_set": ["DataSet"],
@@ -125,7 +132,6 @@
     "saa_s": ["SaaS"],
     "multi_dimensional_dataset": ["MultiDimensionalDataset"],
     "event_store": ["EventStore"],
-    "no_s_q_l": ["NoSQL"],
     "insight": ["Insight"],
     "a_p_i": ["API"],
     "google": ["Google"],
@@ -159,9 +165,12 @@
     "cube_hierarchy": ["CubeHierarchy"],
     "cube_field": ["CubeField"],
     "cube_dimension": ["CubeDimension"],
+    "custom_field": ["CustomField"],
+    "custom_dataset": ["CustomDataset"],
+    "custom_table": ["CustomTable"],
+    "databricks_unity_catalog_tag": ["DatabricksUnityCatalogTag"],
     "kafka": ["Kafka"],
     "azure_service_bus": ["AzureServiceBus"],
-    "cosmos_mongo_d_b": ["CosmosMongoDB"],
     "dynamo_d_b": ["DynamoDB"],
     "mongo_d_b": ["MongoDB"],
     "dbt_tag": ["DbtTag"],
@@ -278,16 +287,11 @@
     "salesforce_organization": ["SalesforceOrganization"],
     "salesforce_dashboard": ["SalesforceDashboard"],
     "salesforce_report": ["SalesforceReport"],
-    "mongo_d_b_collection": ["MongoDBCollection"],
     "dynamo_dbtable": ["DynamoDBTable"],
-    "mongo_d_b_database": ["MongoDBDatabase"],
     "kafka_topic": ["KafkaTopic"],
     "kafka_consumer_group": ["KafkaConsumerGroup"],
     "azure_service_bus_namespace": ["AzureServiceBusNamespace"],
     "azure_service_bus_topic": ["AzureServiceBusTopic"],
-    "cosmos_mongo_d_b_account": ["CosmosMongoDBAccount"],
-    "cosmos_mongo_d_b_collection": ["CosmosMongoDBCollection"],
-    "cosmos_mongo_d_b_database": ["CosmosMongoDBDatabase"],
     "qlik_stream": ["QlikStream"],
     "dynamo_d_b_local_secondary_index": ["DynamoDBLocalSecondaryIndex"],
     "dynamo_d_b_global_secondary_index": ["DynamoDBGlobalSecondaryIndex"],

pyatlan/model/assets/__init__.pyi

@@ -22,8 +22,10 @@ __all__ = [
     "DataQuality",
     "BI",
     "Resource",
+    "Custom",
     "DataMesh",
     "SQL",
+    "NoSQL",
     "Matillion",
     "Dbt",
     "Model",
@@ -59,11 +61,11 @@ __all__ = [
     "TablePartition",
     "Column",
     "SnowflakeStream",
-    "DatabricksUnityCatalogTag",
     "CalculationView",
     "Database",
     "Procedure",
     "SnowflakeTag",
+    "CosmosMongoDB",
     "MatillionGroup",
     "MatillionJob",
     "MatillionProject",
@@ -98,7 +100,12 @@ __all__ = [
     "PowerBIDataflow",
     "PowerBIPage",
     "SnowflakeDynamicTable",
+    "MongoDBCollection",
     "DynamoDBSecondaryIndex",
+    "MongoDBDatabase",
+    "CosmosMongoDBAccount",
+    "CosmosMongoDBCollection",
+    "CosmosMongoDBDatabase",
     "Task",
     "DataSet",
     "TagAttachment",
@@ -155,9 +162,12 @@ __all__ = [
     "CubeHierarchy",
     "CubeField",
     "CubeDimension",
+    "CustomField",
+    "CustomDataset",
+    "CustomTable",
+    "DatabricksUnityCatalogTag",
     "Kafka",
     "AzureServiceBus",
-    "CosmosMongoDB",
     "DynamoDB",
     "MongoDB",
     "DbtTag",
@@ -274,16 +284,11 @@ __all__ = [
     "SalesforceOrganization",
     "SalesforceDashboard",
     "SalesforceReport",
-    "MongoDBCollection",
     "DynamoDBTable",
-    "MongoDBDatabase",
     "KafkaTopic",
     "KafkaConsumerGroup",
     "AzureServiceBusNamespace",
     "AzureServiceBusTopic",
-    "CosmosMongoDBAccount",
-    "CosmosMongoDBCollection",
-    "CosmosMongoDBDatabase",
     "QlikStream",
     "DynamoDBLocalSecondaryIndex",
     "DynamoDBGlobalSecondaryIndex",
@@ -356,13 +361,17 @@ from .core.calculation_view import CalculationView
 from .core.catalog import Catalog
 from .core.column import Column
 from .core.column_process import ColumnProcess
+from .core.cosmos_mongo_d_b import CosmosMongoDB
+from .core.cosmos_mongo_d_b_account import CosmosMongoDBAccount
+from .core.cosmos_mongo_d_b_collection import CosmosMongoDBCollection
+from .core.cosmos_mongo_d_b_database import CosmosMongoDBDatabase
+from .core.custom import Custom
 from .core.data_contract import DataContract
 from .core.data_domain import DataDomain
 from .core.data_mesh import DataMesh
 from .core.data_product import DataProduct
 from .core.data_quality import DataQuality
 from .core.database import Database
-from .core.databricks_unity_catalog_tag import DatabricksUnityCatalogTag
 from .core.dbt import Dbt
 from .core.dbt_metric import DbtMetric
 from .core.dbt_model import DbtModel
@@ -392,8 +401,11 @@ from .core.model_data_model import ModelDataModel
 from .core.model_entity import ModelEntity
 from .core.model_entity_association import ModelEntityAssociation
 from .core.model_version import ModelVersion
+from .core.mongo_d_b_collection import MongoDBCollection
+from .core.mongo_d_b_database import MongoDBDatabase
 from .core.monte_carlo import MonteCarlo
 from .core.namespace import Namespace
+from .core.no_s_q_l import NoSQL
 from .core.persona import Persona
 from .core.power_b_i import PowerBI
 from .core.power_b_i_column import PowerBIColumn
@@ -431,17 +443,17 @@ from .core.table import Table
 from .core.table_partition import TablePartition
 from .core.tag import Tag
 from .core.view import View
-from .cosmos_mongo_d_b import CosmosMongoDB
-from .cosmos_mongo_d_b_account import CosmosMongoDBAccount
-from .cosmos_mongo_d_b_collection import CosmosMongoDBCollection
-from .cosmos_mongo_d_b_database import CosmosMongoDBDatabase
 from .cube import Cube
 from .cube_dimension import CubeDimension
 from .cube_field import CubeField
 from .cube_hierarchy import CubeHierarchy
+from .custom_dataset import CustomDataset
+from .custom_field import CustomField
+from .custom_table import CustomTable
 from .data_set import DataSet
 from .data_studio import DataStudio
 from .data_studio_asset import DataStudioAsset
+from .databricks_unity_catalog_tag import DatabricksUnityCatalogTag
 from .dbt_column_process import DbtColumnProcess
 from .dbt_process import DbtProcess
 from .dbt_tag import DbtTag
@@ -497,10 +509,7 @@ from .mode_query import ModeQuery
 from .mode_report import ModeReport
 from .mode_workspace import ModeWorkspace
 from .mongo_d_b import MongoDB
-from .mongo_d_b_collection import MongoDBCollection
-from .mongo_d_b_database import MongoDBDatabase
 from .multi_dimensional_dataset import MultiDimensionalDataset
-from .no_s_q_l import NoSQL
 from .object_store import ObjectStore
 from .preset import Preset
 from .preset_chart import PresetChart