Cleanup and restructure Pydantic models for Schema

Author: dgarros

infrahub_sdk/checks.py

@@ -17,7 +17,7 @@
     from pathlib import Path
 
     from . import InfrahubClient
-    from .schema import InfrahubCheckDefinitionConfig
+    from .schema.repository import InfrahubCheckDefinitionConfig
 
 INFRAHUB_CHECK_VARIABLE_TO_IMPORT = "INFRAHUB_CHECKS"
 

infrahub_sdk/code_generator.py

@@ -1,17 +1,19 @@
 from collections.abc import Mapping
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import jinja2
 
 from . import protocols as sdk_protocols
 from .ctl.constants import PROTOCOLS_TEMPLATE
 from .schema import (
-    AttributeSchema,
+    AttributeSchemaAPI,
     GenericSchema,
-    MainSchemaTypes,
+    GenericSchemaAPI,
+    MainSchemaTypesAll,
     NodeSchema,
-    ProfileSchema,
-    RelationshipSchema,
+    NodeSchemaAPI,
+    ProfileSchemaAPI,
+    RelationshipSchemaAPI,
 )
 
 ATTRIBUTE_KIND_MAP = {
@@ -40,17 +42,17 @@
 
 
 class CodeGenerator:
-    def __init__(self, schema: dict[str, MainSchemaTypes]):
-        self.generics: dict[str, GenericSchema] = {}
-        self.nodes: dict[str, NodeSchema] = {}
-        self.profiles: dict[str, ProfileSchema] = {}
+    def __init__(self, schema: dict[str, MainSchemaTypesAll]):
+        self.generics: dict[str, Union[GenericSchemaAPI, GenericSchema]] = {}
+        self.nodes: dict[str, Union[NodeSchemaAPI, NodeSchema]] = {}
+        self.profiles: dict[str, ProfileSchemaAPI] = {}
 
         for name, schema_type in schema.items():
-            if isinstance(schema_type, GenericSchema):
+            if isinstance(schema_type, (GenericSchemaAPI, GenericSchema)):
                 self.generics[name] = schema_type
-            if isinstance(schema_type, NodeSchema):
+            if isinstance(schema_type, (NodeSchemaAPI, NodeSchema)):
                 self.nodes[name] = schema_type
-            if isinstance(schema_type, ProfileSchema):
+            if isinstance(schema_type, ProfileSchemaAPI):
                 self.profiles[name] = schema_type
 
         self.base_protocols = [
@@ -92,7 +94,7 @@ def _jinja2_filter_inheritance(value: dict[str, Any]) -> str:
         return ", ".join(inherit_from)
 
     @staticmethod
-    def _jinja2_filter_render_attribute(value: AttributeSchema) -> str:
+    def _jinja2_filter_render_attribute(value: AttributeSchemaAPI) -> str:
         attribute_kind: str = ATTRIBUTE_KIND_MAP[value.kind]
 
         if value.optional:
@@ -101,7 +103,7 @@ def _jinja2_filter_render_attribute(value: AttributeSchema) -> str:
         return f"{value.name}: {attribute_kind}"
 
     @staticmethod
-    def _jinja2_filter_render_relationship(value: RelationshipSchema, sync: bool = False) -> str:
+    def _jinja2_filter_render_relationship(value: RelationshipSchemaAPI, sync: bool = False) -> str:
         name = value.name
         cardinality = value.cardinality
 
@@ -116,12 +118,12 @@ def _jinja2_filter_render_relationship(value: RelationshipSchema, sync: bool = F
 
     @staticmethod
     def _sort_and_filter_models(
-        models: Mapping[str, MainSchemaTypes], filters: Optional[list[str]] = None
-    ) -> list[MainSchemaTypes]:
+        models: Mapping[str, MainSchemaTypesAll], filters: Optional[list[str]] = None
+    ) -> list[MainSchemaTypesAll]:
         if filters is None:
             filters = ["CoreNode"]
 
-        filtered: list[MainSchemaTypes] = []
+        filtered: list[MainSchemaTypesAll] = []
         for name, model in models.items():
             if name in filters:
                 continue

infrahub_sdk/ctl/check.py

@@ -17,7 +17,7 @@
 from ..ctl.repository import get_repository_config
 from ..ctl.utils import catch_exception, execute_graphql_query
 from ..exceptions import ModuleImportError
-from ..schema import InfrahubCheckDefinitionConfig, InfrahubRepositoryConfig
+from ..schema.repository import InfrahubCheckDefinitionConfig, InfrahubRepositoryConfig
 
 app = typer.Typer()
 console = Console()

infrahub_sdk/ctl/cli_commands.py

@@ -38,11 +38,8 @@
 from ..ctl.validate import app as validate_app
 from ..exceptions import GraphQLError, ModuleImportError
 from ..jinja2 import identify_faulty_jinja_code
-from ..schema import (
-    InfrahubRepositoryConfig,
-    MainSchemaTypes,
-    SchemaRoot,
-)
+from ..schema import MainSchemaTypesAll, SchemaRoot
+from ..schema.repository import InfrahubRepositoryConfig
 from ..utils import get_branch, write_to_file
 from ..yaml import SchemaFile
 from .exporter import dump
@@ -364,7 +361,7 @@ def protocols(
 ) -> None:
     """Export Python protocols corresponding to a schema."""
 
-    schema: dict[str, MainSchemaTypes] = {}
+    schema: dict[str, MainSchemaTypesAll] = {}
 
     if schemas:
         schemas_data = load_yamlfile_from_disk_and_exit(paths=schemas, file_type=SchemaFile, console=console)

infrahub_sdk/node.py

@@ -14,15 +14,15 @@
     UninitializedError,
 )
 from .graphql import Mutation, Query
-from .schema import GenericSchema, RelationshipCardinality, RelationshipKind
+from .schema import GenericSchemaAPI, RelationshipCardinality, RelationshipKind
 from .utils import compare_lists, get_flat_value
 from .uuidt import UUIDT
 
 if TYPE_CHECKING:
     from typing_extensions import Self
 
     from .client import InfrahubClient, InfrahubClientSync
-    from .schema import AttributeSchema, MainSchemaTypes, RelationshipSchema
+    from .schema import AttributeSchemaAPI, MainSchemaTypesAPI, RelationshipSchemaAPI
 
 # pylint: disable=too-many-lines
 
@@ -46,7 +46,7 @@
 class Attribute:
     """Represents an attribute of a Node, including its schema, value, and properties."""
 
-    def __init__(self, name: str, schema: AttributeSchema, data: Union[Any, dict]):
+    def __init__(self, name: str, schema: AttributeSchemaAPI, data: Union[Any, dict]):
         """
         Args:
             name (str): The name of the attribute.
@@ -143,7 +143,7 @@ def _generate_mutation_query(self) -> dict[str, Any]:
 class RelatedNodeBase:
     """Base class for representing a related node in a relationship."""
 
-    def __init__(self, branch: str, schema: RelationshipSchema, data: Union[Any, dict], name: Optional[str] = None):
+    def __init__(self, branch: str, schema: RelationshipSchemaAPI, data: Union[Any, dict], name: Optional[str] = None):
         """
         Args:
             branch (str): The branch where the related node resides.
@@ -300,7 +300,7 @@ def __init__(
         self,
         client: InfrahubClient,
         branch: str,
-        schema: RelationshipSchema,
+        schema: RelationshipSchemaAPI,
         data: Union[Any, dict],
         name: Optional[str] = None,
     ):
@@ -347,7 +347,7 @@ def __init__(
         self,
         client: InfrahubClientSync,
         branch: str,
-        schema: RelationshipSchema,
+        schema: RelationshipSchemaAPI,
         data: Union[Any, dict],
         name: Optional[str] = None,
     ):
@@ -390,7 +390,7 @@ def get(self) -> InfrahubNodeSync:
 class RelationshipManagerBase:
     """Base class for RelationshipManager and RelationshipManagerSync"""
 
-    def __init__(self, name: str, branch: str, schema: RelationshipSchema):
+    def __init__(self, name: str, branch: str, schema: RelationshipSchemaAPI):
         """
         Args:
             name (str): The name of the relationship.
@@ -473,7 +473,7 @@ def __init__(
         client: InfrahubClient,
         node: InfrahubNode,
         branch: str,
-        schema: RelationshipSchema,
+        schema: RelationshipSchemaAPI,
         data: Union[Any, dict],
     ):
         """
@@ -568,7 +568,7 @@ def __init__(
         client: InfrahubClientSync,
         node: InfrahubNodeSync,
         branch: str,
-        schema: RelationshipSchema,
+        schema: RelationshipSchemaAPI,
         data: Union[Any, dict],
     ):
         """
@@ -657,12 +657,12 @@ def remove(self, data: Union[str, RelatedNodeSync, dict]) -> None:
 class InfrahubNodeBase:
     """Base class for InfrahubNode and InfrahubNodeSync"""
 
-    def __init__(self, schema: MainSchemaTypes, branch: str, data: Optional[dict] = None) -> None:
+    def __init__(self, schema: MainSchemaTypesAPI, branch: str, data: Optional[dict] = None) -> None:
         """
         Args:
-            schema (MainSchemaTypes): The schema of the node.
-            branch (str): The branch where the node resides.
-            data (Optional[dict]): Optional data to initialize the node.
+            schema: The schema of the node.
+            branch: The branch where the node resides.
+            data: Optional data to initialize the node.
         """
         self._schema = schema
         self._data = data
@@ -1035,16 +1035,16 @@ class InfrahubNode(InfrahubNodeBase):
     def __init__(
         self,
         client: InfrahubClient,
-        schema: MainSchemaTypes,
+        schema: MainSchemaTypesAPI,
         branch: Optional[str] = None,
         data: Optional[dict] = None,
     ) -> None:
         """
         Args:
-            client (InfrahubClient): The client used to interact with the backend.
-            schema (MainSchemaTypes): The schema of the node.
-            branch (Optional[str]): The branch where the node resides.
-            data (Optional[dict]): Optional data to initialize the node.
+            client: The client used to interact with the backend.
+            schema: The schema of the node.
+            branch: The branch where the node resides.
+            data: Optional data to initialize the node.
         """
         self._client = client
         self.__class__ = type(f"{schema.kind}InfrahubNode", (self.__class__,), {})
@@ -1060,7 +1060,7 @@ async def from_graphql(
         client: InfrahubClient,
         branch: str,
         data: dict,
-        schema: Optional[MainSchemaTypes] = None,
+        schema: Optional[MainSchemaTypesAPI] = None,
         timeout: Optional[int] = None,
     ) -> Self:
         if not schema:
@@ -1146,7 +1146,7 @@ async def save(
         if update_group_context is None and self._client.mode == InfrahubClientMode.TRACKING:
             update_group_context = True
 
-        if not isinstance(self._schema, GenericSchema):
+        if not isinstance(self._schema, GenericSchemaAPI):
             if "CoreGroup" in self._schema.inherit_from:
                 await self._client.group_context.add_related_groups(
                     ids=[self.id], update_group_context=update_group_context
@@ -1183,7 +1183,7 @@ async def generate_query_data(
             )
         )
 
-        if isinstance(self._schema, GenericSchema) and fragment:
+        if isinstance(self._schema, GenericSchemaAPI) and fragment:
             for child in self._schema.used_by:
                 child_schema = await self._client.schema.get(kind=child)
                 child_node = InfrahubNode(client=self._client, schema=child_schema)
@@ -1540,7 +1540,7 @@ class InfrahubNodeSync(InfrahubNodeBase):
     def __init__(
         self,
         client: InfrahubClientSync,
-        schema: MainSchemaTypes,
+        schema: MainSchemaTypesAPI,
         branch: Optional[str] = None,
         data: Optional[dict] = None,
     ) -> None:
@@ -1565,7 +1565,7 @@ def from_graphql(
         client: InfrahubClientSync,
         branch: str,
         data: dict,
-        schema: Optional[MainSchemaTypes] = None,
+        schema: Optional[MainSchemaTypesAPI] = None,
         timeout: Optional[int] = None,
     ) -> Self:
         if not schema:
@@ -1648,7 +1648,7 @@ def save(
         if update_group_context is None and self._client.mode == InfrahubClientMode.TRACKING:
             update_group_context = True
 
-        if not isinstance(self._schema, GenericSchema):
+        if not isinstance(self._schema, GenericSchemaAPI):
             if "CoreGroup" in self._schema.inherit_from:
                 self._client.group_context.add_related_groups(ids=[self.id], update_group_context=update_group_context)
             else:
@@ -1681,7 +1681,7 @@ def generate_query_data(
             )
         )
 
-        if isinstance(self._schema, GenericSchema) and fragment:
+        if isinstance(self._schema, GenericSchemaAPI) and fragment:
             for child in self._schema.used_by:
                 child_schema = self._client.schema.get(kind=child)
                 child_node = InfrahubNodeSync(client=self._client, schema=child_schema)

infrahub_sdk/pytest_plugin/items/base.py

@@ -13,7 +13,7 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from ...schema import InfrahubRepositoryConfigElement
+    from ...schema.repository import InfrahubRepositoryConfigElement
     from ..models import InfrahubTest
 
 

infrahub_sdk/pytest_plugin/items/check.py

@@ -15,7 +15,7 @@
     from pytest import ExceptionInfo
 
     from ...checks import InfrahubCheck
-    from ...schema import InfrahubRepositoryConfigElement
+    from ...schema.repository import InfrahubRepositoryConfigElement
     from ..models import InfrahubTest
 
 

infrahub_sdk/pytest_plugin/items/python_transform.py

@@ -14,7 +14,7 @@
 if TYPE_CHECKING:
     from pytest import ExceptionInfo
 
-    from ...schema import InfrahubRepositoryConfigElement
+    from ...schema.repository import InfrahubRepositoryConfigElement
     from ...transforms import InfrahubTransform
     from ..models import InfrahubTest
 

infrahub_sdk/query_groups.py

@@ -9,7 +9,7 @@
 if TYPE_CHECKING:
     from .client import InfrahubClient, InfrahubClientSync
     from .node import InfrahubNode, InfrahubNodeSync, RelatedNodeBase
-    from .schema import MainSchemaTypes
+    from .schema import MainSchemaTypesAPI
 
 
 class InfrahubGroupContextBase:
@@ -63,7 +63,7 @@ def _generate_group_name(self, suffix: Optional[str] = None) -> str:
 
         return group_name
 
-    def _generate_group_description(self, schema: MainSchemaTypes) -> str:
+    def _generate_group_description(self, schema: MainSchemaTypesAPI) -> str:
         """Generate the description of the group from the params
         and ensure it's not longer than the maximum length of the description field."""
         if not self.params: