[generator] Fixed template imports

Author: Aryamanz29

pyatlan/generator/templates/imports.jinja2

@@ -120,6 +120,7 @@ from pyatlan.model.structs import (
 from pyatlan.utils import (
     init_guid,
     next_id,
+    to_camel_case,
     validate_required_fields,
     validate_single_required_field,
     get_parent_qualified_name,

pyatlan/generator/templates/module.jinja2

@@ -20,6 +20,11 @@ from pyatlan.model.enums import (
 from .asset import SelfAsset
 {% endif %}
 
+{% if asset_info.name == 'AIModel' %}
+from .asset import Asset
+from .process import Process
+{% endif %}
+
 {% set entity_def = asset_info.entity_def %}
 
 {% set file_name = 'methods/imports/' + entity_def.name | to_snake_case + '.jinja2' %}

pyatlan/model/assets/__init__.py

@@ -88,8 +88,8 @@
         "SnowflakeStage",
         "DatabricksUnityCatalogTag",
         "SnowflakeStream",
-        "Database",
         "CalculationView",
+        "Database",
         "Procedure",
         "SnowflakeTag",
         "MatillionGroup",

pyatlan/model/assets/__init__.pyi

@@ -85,8 +85,8 @@ __all__ = [
     "SnowflakeStage",
     "DatabricksUnityCatalogTag",
     "SnowflakeStream",
-    "Database",
     "CalculationView",
+    "Database",
     "Procedure",
     "SnowflakeTag",
     "MatillionGroup",

pyatlan/model/assets/core/__init__.py

@@ -219,8 +219,8 @@
 SnowflakeStage.Attributes.update_forward_refs(**localns)
 DatabricksUnityCatalogTag.Attributes.update_forward_refs(**localns)
 SnowflakeStream.Attributes.update_forward_refs(**localns)
-Database.Attributes.update_forward_refs(**localns)
 CalculationView.Attributes.update_forward_refs(**localns)
+Database.Attributes.update_forward_refs(**localns)
 Procedure.Attributes.update_forward_refs(**localns)
 SnowflakeTag.Attributes.update_forward_refs(**localns)
 MatillionGroup.Attributes.update_forward_refs(**localns)

pyatlan/model/structs.py

@@ -45,19 +45,6 @@ def flatten_structs_attributes(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         return values
 
 
-class MCRuleSchedule(AtlanObject):
-    """Description"""
-
-    mc_rule_schedule_type: Optional[str] = Field(default=None, description="")
-    mc_rule_schedule_interval_in_minutes: Optional[int] = Field(
-        default=None, description=""
-    )
-    mc_rule_schedule_start_time: Optional[datetime] = Field(
-        default=None, description=""
-    )
-    mc_rule_schedule_crontab: Optional[str] = Field(default=None, description="")
-
-
 class DbtJobRun(AtlanObject):
     """Description"""
 
@@ -74,18 +61,24 @@ class DbtJobRun(AtlanObject):
     dbt_compiled_code: Optional[str] = Field(default=None, description="")
 
 
-class AwsCloudWatchMetric(AtlanObject):
+class MCRuleSchedule(AtlanObject):
     """Description"""
 
-    aws_cloud_watch_metric_name: str = Field(description="")
-    aws_cloud_watch_metric_scope: str = Field(description="")
+    mc_rule_schedule_type: Optional[str] = Field(default=None, description="")
+    mc_rule_schedule_interval_in_minutes: Optional[int] = Field(
+        default=None, description=""
+    )
+    mc_rule_schedule_start_time: Optional[datetime] = Field(
+        default=None, description=""
+    )
+    mc_rule_schedule_crontab: Optional[str] = Field(default=None, description="")
 
 
-class Histogram(AtlanObject):
+class AwsCloudWatchMetric(AtlanObject):
     """Description"""
 
-    boundaries: Set[float] = Field(description="")
-    frequencies: Set[float] = Field(description="")
+    aws_cloud_watch_metric_name: str = Field(description="")
+    aws_cloud_watch_metric_scope: str = Field(description="")
 
 
 class Action(AtlanObject):
@@ -97,18 +90,18 @@ class Action(AtlanObject):
     task_action_display_text: Optional[str] = Field(default=None, description="")
 
 
-class ColumnValueFrequencyMap(AtlanObject):
+class Histogram(AtlanObject):
     """Description"""
 
-    column_value: Optional[str] = Field(default=None, description="")
-    column_value_frequency: Optional[int] = Field(default=None, description="")
+    boundaries: Set[float] = Field(description="")
+    frequencies: Set[float] = Field(description="")
 
 
-class SourceTagAttachmentValue(AtlanObject):
+class ColumnValueFrequencyMap(AtlanObject):
     """Description"""
 
-    tag_attachment_key: Optional[str] = Field(default=None, description="")
-    tag_attachment_value: Optional[str] = Field(default=None, description="")
+    column_value: Optional[str] = Field(default=None, description="")
+    column_value_frequency: Optional[int] = Field(default=None, description="")
 
 
 class BadgeCondition(AtlanObject):
@@ -143,6 +136,13 @@ def create(
     badge_condition_colorhex: Optional[str] = Field(default=None, description="")
 
 
+class SourceTagAttachmentValue(AtlanObject):
+    """Description"""
+
+    tag_attachment_key: Optional[str] = Field(default=None, description="")
+    tag_attachment_value: Optional[str] = Field(default=None, description="")
+
+
 class StarredDetails(AtlanObject):
     """Description"""
 
@@ -461,22 +461,22 @@ class SourceTagAttribute(AtlanObject):
     )
 
 
-MCRuleSchedule.update_forward_refs()
-
 DbtJobRun.update_forward_refs()
 
-AwsCloudWatchMetric.update_forward_refs()
+MCRuleSchedule.update_forward_refs()
 
-Histogram.update_forward_refs()
+AwsCloudWatchMetric.update_forward_refs()
 
 Action.update_forward_refs()
 
-ColumnValueFrequencyMap.update_forward_refs()
+Histogram.update_forward_refs()
 
-SourceTagAttachmentValue.update_forward_refs()
+ColumnValueFrequencyMap.update_forward_refs()
 
 BadgeCondition.update_forward_refs()
 
+SourceTagAttachmentValue.update_forward_refs()
+
 StarredDetails.update_forward_refs()
 
 AwsTag.update_forward_refs()