Migrate to pydantic v2

pyatlan/client/admin.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
 
-from pydantic import ValidationError, parse_obj_as, validate_arguments
+from pydantic.v1 import ValidationError, parse_obj_as, validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import ADMIN_EVENTS, KEYCLOAK_EVENTS

pyatlan/client/asset.py

@@ -11,7 +11,7 @@
 from warnings import warn
 
 import requests
-from pydantic import (
+from pydantic.v1 import (
     StrictStr,
     ValidationError,
     constr,

pyatlan/client/atlan.py

@@ -16,7 +16,7 @@
 from warnings import warn
 
 import requests
-from pydantic import (
+from pydantic.v1 import (
     BaseSettings,
     HttpUrl,
     PrivateAttr,

pyatlan/client/audit.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2022 Atlan Pte. Ltd.
-from pydantic import ValidationError, parse_obj_as, validate_arguments
+from pydantic.v1 import ValidationError, parse_obj_as, validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import AUDIT_SEARCH

pyatlan/client/credential.py

@@ -1,4 +1,4 @@
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (

pyatlan/client/group.py

@@ -2,7 +2,7 @@
 # Copyright 2022 Atlan Pte. Ltd.
 from typing import Optional
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (

pyatlan/client/query.py

@@ -1,4 +1,4 @@
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import RUN_QUERY

pyatlan/client/role.py

@@ -2,7 +2,7 @@
 # Copyright 2022 Atlan Pte. Ltd.
 from typing import Optional
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import GET_ROLES

pyatlan/client/search_log.py

@@ -1,6 +1,6 @@
 from typing import Union
 
-from pydantic import ValidationError, parse_obj_as, validate_arguments
+from pydantic.v1 import ValidationError, parse_obj_as, validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import SEARCH_LOG

pyatlan/client/token.py

@@ -4,7 +4,7 @@
 
 from typing import Optional
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import DELETE_API_TOKEN, GET_API_TOKENS, UPSERT_API_TOKEN

pyatlan/client/typedef.py

@@ -2,7 +2,7 @@
 # Copyright 2022 Atlan Pte. Ltd.
 from typing import Union
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (
@@ -32,7 +32,7 @@ def _build_typedef_request(typedef: TypeDef) -> TypeDefResponse:
             entity_defs=[],
             relationship_defs=[],
             custom_metadata_defs=[],
-        )
+        )  # type: ignore[call-arg]
     elif isinstance(typedef, CustomMetadataDef):
         # Set up the request payload...
         payload = TypeDefResponse(
@@ -42,7 +42,7 @@ def _build_typedef_request(typedef: TypeDef) -> TypeDefResponse:
             entity_defs=[],
             relationship_defs=[],
             custom_metadata_defs=[typedef],
-        )
+        )  # type: ignore[call-arg]
     elif isinstance(typedef, EnumDef):
         # Set up the request payload...
         payload = TypeDefResponse(
@@ -52,7 +52,7 @@ def _build_typedef_request(typedef: TypeDef) -> TypeDefResponse:
             entity_defs=[],
             relationship_defs=[],
             custom_metadata_defs=[],
-        )
+        )  # type: ignore[call-arg]
     else:
         raise ErrorCode.UNABLE_TO_UPDATE_TYPEDEF_CATEGORY.exception_with_parameters(
             typedef.category.value

pyatlan/client/user.py

@@ -4,7 +4,7 @@
 
 from typing import Any, Optional
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (
@@ -59,7 +59,7 @@ def create(
         cur = CreateUserRequest(users=[])
         for user in users:
             role_name = str(user.workspace_role)
-            if role_id := RoleCache.get_id_for_name(role_name):
+            if role_id := RoleCache.get_id_for_name(role_name) and user.email:
                 to_create = CreateUserRequest.CreateUser(
                     email=user.email,
                     role_name=role_name,
@@ -342,7 +342,7 @@ def _add_as(
         :raises NotFoundError: if the asset to which to add the API token as a viewer cannot be found
         """
         from pyatlan.client.atlan import client_connection
-        from pyatlan.model.assets.asset00 import Asset
+        from pyatlan.model.assets import Asset
         from pyatlan.model.fluent_search import FluentSearch
 
         if keyword_field not in [Asset.ADMIN_USERS, Asset.VIEWER_USERS]:

pyatlan/client/workflow.py

@@ -4,7 +4,7 @@
 from time import sleep
 from typing import Optional, Union, overload
 
-from pydantic import validate_arguments
+from pydantic.v1 import validate_arguments
 
 from pyatlan.client.common import ApiCaller
 from pyatlan.client.constants import (

pyatlan/generator/class_generator.py

@@ -112,6 +112,7 @@ class ModuleInfo:
     count: int = 0
     modules: set["ModuleInfo"] = set()
     modules_by_asset_name: dict[str, str] = {}
+    assets: dict[str, "AssetInfo"] = {}
 
     @classmethod
     def check_for_circular_module_dependencies(cls):
@@ -210,6 +211,8 @@ def __init__(self, name: str, entity_def: EntityDef):
         self.required_asset_infos: set["AssetInfo"] = set()
         self.circular_dependencies: set["AssetInfo"] = set()
         self.order: int = 0
+        self.module_name = to_snake_case(name)
+        self.super_type: Optional[AssetInfo] = None
 
     def __hash__(self):
         return hash(self._name)
@@ -228,6 +231,45 @@ def super_class(self):
         else:
             return self.entity_def.super_types[0]
 
+    @property
+    def import_super_class(self):
+        if self._name == REFERENCEABLE:
+            return ""
+        super_type = AssetInfo.asset_info_by_name[self.entity_def.super_types[0]]
+        return f"from .{super_type.module_name} import {super_type.name}"
+
+    @property
+    def get_update_forward_refs(self):
+        import_set = set()
+        asset_name = self.name
+        current_name = self.name
+        while current_name != "Referenceable" and self.hierarchy_graph.predecessors(
+            current_name
+        ):
+            parent_asset_name = next(
+                iter(self.hierarchy_graph.predecessors(current_name))
+            )
+            if parent_asset_name == "Referenceable":
+                break
+            parent_asset = self.asset_info_by_name[
+                parent_asset_name
+            ].required_asset_infos
+
+            for parent_module in parent_asset:
+                if asset_name != parent_module.name:
+                    import_set.add(parent_module.name)
+
+            current_name = parent_asset_name
+
+        return {key: key for key in import_set}
+
+    @property
+    def imports_for_referenced_assets(self):
+        return [
+            f"from .{required_asset.module_name} import {required_asset.name} # noqa"
+            for required_asset in self.required_asset_infos
+        ]
+
     def update_attribute_defs(self):
         def get_ancestor_relationship_defs(
             ancestor_name: str, ancestor_relationship_defs
@@ -284,10 +326,6 @@ def update_required_asset_names(self) -> None:
             for a in relationship_attribute_defs
             if a["name"] not in attributes_to_remove
         ]
-        if self.entity_def.super_types:
-            self.required_asset_infos.add(
-                AssetInfo.asset_info_by_name[self.entity_def.super_types[0]]
-            )
 
     def merge_attributes(self, entity_def):
         def merge_them(s, a):
@@ -357,10 +395,8 @@ def create_modules(cls):
                 asset_info = cls.asset_info_by_name[asset_name]
                 asset_info.order = order
                 order += 1
-                if asset_info.module_info is None:
-                    ModuleInfo(asset_info=asset_info)
-                else:
-                    asset_info.module_info.add_asset_info(asset_info=asset_info)
+
+                ModuleInfo.assets[asset_name] = asset_info
 
 
 class AttributeType(Enum):
@@ -612,40 +648,32 @@ def get_ancestor_relationship_defs(
         )
 
     def render_modules(self, modules: list[ModuleInfo]):
-        for module in modules:
-            self.render_module(module)
-        self.render_init(modules)
+        self.render_init(modules)  # type: ignore
 
-    def render_module(self, module: ModuleInfo):
+    def render_module(self, asset_info: AssetInfo):
         template = self.environment.get_template("module.jinja2")
         content = template.render(
             {
-                "module": module,
+                "asset_info": asset_info,
                 "existz": os.path.exists,
-                "module_name": module.name,
-                "modules_by_asset_name": ModuleInfo.modules_by_asset_name,
             }
         )
-        with (ASSETS_DIR / f"{module.name}.py").open("w") as script:
+        with (ASSETS_DIR / f"{asset_info.module_name}.py").open("w") as script:
             script.write(content)
 
-    def render_init(self, modules: list[ModuleInfo]):
-        imports = sorted(
-            {
-                f"from .{asset_info.module_info.name} import {asset_info.name}"
-                for module in modules
-                if module.status == ModuleStatus.ACTIVE
-                for asset_info in module.asset_infos
-                if asset_info.module_info
-            }
-        )
+    def render_init(self, assets: list[AssetInfo]):
+        asset_names = [asset.name for asset in assets]
+        asset_imports = [
+            f"from .{asset.module_name} import {asset.name}" for asset in assets
+        ]
+
         template = self.environment.get_template("init.jinja2")
         content = template.render(
-            {
-                "imports": imports,
-            }
+            {"asset_imports": asset_imports, "asset_names": asset_names}
         )
-        with (ASSETS_DIR / "__init__.py").open("w") as script:
+
+        init_path = ASSETS_DIR / "__init__.py"
+        with init_path.open("w") as script:
             script.write(content)
 
     def render_structs(self, struct_defs):
@@ -788,13 +816,12 @@ def create(cls, enum_defs):
     AssetInfo.set_entity_defs(type_defs.entity_defs)
     AssetInfo.update_all_circular_dependencies()
     AssetInfo.create_modules()
-    ModuleInfo.check_for_circular_module_dependencies()
     for file in (ASSETS_DIR).glob("*.py"):
         file.unlink()
     generator = Generator()
-    generator.render_modules(
-        [m for m in ModuleInfo.modules if m.status == ModuleStatus.ACTIVE]
-    )
+    for asset_info in ModuleInfo.assets.values():
+        generator.render_module(asset_info)
+    generator.render_init(ModuleInfo.assets.values())  # type: ignore
     generator.render_structs(type_defs.struct_defs)
     EnumDefInfo.create(type_defs.enum_defs)
     generator.render_enums(EnumDefInfo.enum_def_info)

pyatlan/generator/templates/imports.jinja2

@@ -8,7 +8,7 @@ from io import StringIO
 from typing import Any, ClassVar, Optional, Set, Type, TypeVar, TYPE_CHECKING, cast, overload
 from urllib.parse import quote, unquote
 
-from pydantic import Field, PrivateAttr, StrictStr, root_validator, validator
+from pydantic.v1 import Field, PrivateAttr, StrictStr, root_validator, validator
 
 from pyatlan.errors import ErrorCode
 from pyatlan.model.core import Announcement, AtlanObject, AtlanTag, AtlanTagName, Meaning
@@ -103,5 +103,6 @@ from pyatlan.utils import (
     next_id,
     to_camel_case,
     validate_required_fields,
+    validate_single_required_field,
     get_parent_qualified_name,
 )

pyatlan/generator/templates/init.jinja2

@@ -1,5 +1,11 @@
 # Copyright 2022 Atlan Pte. Ltd.
-from .asset00 import validate_single_required_field
-{%  for import in imports -%}
-{{ import }}
+# isort: skip_file
+{% for asset_import in asset_imports -%}
+{{ asset_import }}
 {% endfor %}
+
+# Update asset forward references:
+localns = locals()
+{%- for asset_name in asset_names %}
+{{ asset_name }}.Attributes.update_forward_refs(**localns)
+{%- endfor %}

pyatlan/generator/templates/methods/asset/asset.jinja2

@@ -1,4 +1,5 @@
 
+
     _subtypes_:dict[str, type] = dict()
 
     def __init_subclass__(cls, type_name=None):

pyatlan/generator/templates/methods/asset/purpose.jinja2

@@ -2,7 +2,7 @@
     @classmethod
     # @validate_arguments()
     @init_guid
-    def create(cls, *, name: str, atlan_tags: list[str]) -> {{ entity_def.name }}:
+    def create(cls, *, name: str, atlan_tags: list[AtlanTagName]) -> {{ entity_def.name }}:
         validate_required_fields(["name", "atlan_tags"], [name, atlan_tags])
         attributes = Purpose.Attributes.create(name=name, atlan_tags=atlan_tags)
         return cls(attributes=attributes)

pyatlan/generator/templates/methods/attribute/a_p_i_path.jinja2

@@ -24,5 +24,5 @@
                 connection_qualified_name=f"{fields[0]}/{fields[1]}/{fields[2]}",
                 qualified_name=f"{spec_qualified_name}{path_raw_uri}",
                 connector_name=connector_type.value,
-                apiSpec=APISpec.ref_by_qualified_name(spec_qualified_name),
+                api_spec=APISpec.ref_by_qualified_name(spec_qualified_name),
             )

pyatlan/generator/templates/methods/attribute/atlas_glossary.jinja2

@@ -5,4 +5,4 @@
         def create(cls, *, name: StrictStr, icon: Optional[AtlanIcon] = None)->{{ entity_def.name }}.Attributes:
             validate_required_fields(["name"], [name])
             icon_str = icon.value if icon is not None else None
-            return AtlasGlossary.Attributes(name=name, qualified_name=next_id(), icon=icon_str)
+            return AtlasGlossary.Attributes(name=name, qualified_name=next_id(), asset_icon=icon_str)

pyatlan/generator/templates/methods/attribute/data_domain.jinja2

@@ -24,5 +24,5 @@
                 name=name,
                 parent_domain=parent_domain,
                 qualified_name=qualified_name,
-                icon=icon_str,
+                asset_icon=icon_str,
             )

pyatlan/generator/templates/methods/attribute/data_product.jinja2

@@ -25,5 +25,5 @@
                 data_product_assets_d_s_l=assets_dsl,
                 data_domain=domain,
                 qualified_name=f"default/product/{camel_case_name}",
-                icon=icon_str,
+                asset_icon=icon_str,
             )

pyatlan/generator/templates/methods/attribute/purpose.jinja2

@@ -2,7 +2,7 @@
         @classmethod
         # @validate_arguments()
         @init_guid
-        def create(cls, name: str, atlan_tags: list[str]) -> {{ entity_def.name }}.Attributes:
+        def create(cls, name: str, atlan_tags: list[AtlanTagName]) -> {{ entity_def.name }}.Attributes:
             validate_required_fields(["name", "atlan_tags"], [name, atlan_tags])
             return Purpose.Attributes(
                 qualified_name=name,

pyatlan/generator/templates/methods/imports/auth_policy.jinja2

@@ -1 +1 @@
-from .{{ modules_by_asset_name['Asset'] }} import SelfAsset
+from .asset import SelfAsset