refactor(schema): name generated models InfrahubSchema{Write,Read} + per-family Write/Read

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: dgarros

infrahub_sdk/schema/__init__.py

@@ -24,6 +24,8 @@
 from ..protocols_base import CoreNodeBase
 from ..queries import SCHEMA_HASH_SYNC_STATUS
 from .export import RESTRICTED_NAMESPACES, NamespaceExport, SchemaExport, schema_to_export_dict
+from .generated.read import InfrahubSchemaRead
+from .generated.write import InfrahubSchemaWrite
 from .main import (
     AttributeSchema,
     AttributeSchemaAPI,
@@ -57,6 +59,8 @@
     "BranchSupportType",
     "GenericSchema",
     "GenericSchemaAPI",
+    "InfrahubSchemaRead",
+    "InfrahubSchemaWrite",
     "NamespaceExport",
     "NodeSchema",
     "NodeSchemaAPI",

infrahub_sdk/schema/generated/read.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 
-class GeneratedAttributeSchema(BaseModel):
+class AttributeSchemaRead(BaseModel):
     model_config = ConfigDict()
     id: str | None = Field(
         default=None,
@@ -132,7 +132,7 @@ class GeneratedAttributeSchema(BaseModel):
     )
 
 
-class GeneratedRelationshipSchema(BaseModel):
+class RelationshipSchemaRead(BaseModel):
     model_config = ConfigDict()
     id: str | None = Field(
         default=None,
@@ -241,7 +241,7 @@ class GeneratedRelationshipSchema(BaseModel):
     )
 
 
-class GeneratedBaseNodeSchema(BaseModel):
+class BaseNodeSchemaRead(BaseModel):
     model_config = ConfigDict()
     id: str | None = Field(
         default=None,
@@ -320,17 +320,17 @@ class GeneratedBaseNodeSchema(BaseModel):
         default="present",
         description="Expected state of the node/generic after loading the schema",
     )
-    attributes: list[GeneratedAttributeSchema] = Field(
+    attributes: list[AttributeSchemaRead] = Field(
         default_factory=list,
         description="Node attributes",
     )
-    relationships: list[GeneratedRelationshipSchema] = Field(
+    relationships: list[RelationshipSchemaRead] = Field(
         default_factory=list,
         description="Node Relationships",
     )
 
 
-class GeneratedNodeSchema(GeneratedBaseNodeSchema):
+class NodeSchemaRead(BaseNodeSchemaRead):
     model_config = ConfigDict()
     inherit_from: list[str] = Field(
         default_factory=list,
@@ -358,7 +358,7 @@ class GeneratedNodeSchema(GeneratedBaseNodeSchema):
     )
 
 
-class GeneratedGenericSchema(GeneratedBaseNodeSchema):
+class GenericSchemaRead(BaseNodeSchemaRead):
     model_config = ConfigDict()
     hierarchical: bool = Field(
         default=False,
@@ -376,3 +376,8 @@ class GeneratedGenericSchema(GeneratedBaseNodeSchema):
         default=None,
         description="Nodes inheriting from this Generic schema must belong to one of the listed namespaces",
     )
+
+
+class InfrahubSchemaRead(BaseModel):
+    nodes: list[NodeSchemaRead] = Field(default_factory=list)
+    generics: list[GenericSchemaRead] = Field(default_factory=list)

infrahub_sdk/schema/generated/write.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 
-class GeneratedAttributeSchema(BaseModel):
+class AttributeSchemaWrite(BaseModel):
     model_config = ConfigDict(extra="forbid")
     id: str | None = Field(
         default=None,
@@ -128,7 +128,7 @@ class GeneratedAttributeSchema(BaseModel):
     )
 
 
-class GeneratedRelationshipSchema(BaseModel):
+class RelationshipSchemaWrite(BaseModel):
     model_config = ConfigDict(extra="forbid")
     id: str | None = Field(
         default=None,
@@ -229,7 +229,7 @@ class GeneratedRelationshipSchema(BaseModel):
     )
 
 
-class GeneratedBaseNodeSchema(BaseModel):
+class BaseNodeSchemaWrite(BaseModel):
     model_config = ConfigDict(extra="forbid")
     id: str | None = Field(
         default=None,
@@ -308,17 +308,17 @@ class GeneratedBaseNodeSchema(BaseModel):
         default="present",
         description="Expected state of the node/generic after loading the schema",
     )
-    attributes: list[GeneratedAttributeSchema] = Field(
+    attributes: list[AttributeSchemaWrite] = Field(
         default_factory=list,
         description="Node attributes",
     )
-    relationships: list[GeneratedRelationshipSchema] = Field(
+    relationships: list[RelationshipSchemaWrite] = Field(
         default_factory=list,
         description="Node Relationships",
     )
 
 
-class GeneratedNodeSchema(GeneratedBaseNodeSchema):
+class NodeSchemaWrite(BaseNodeSchemaWrite):
     model_config = ConfigDict(extra="forbid")
     inherit_from: list[str] = Field(
         default_factory=list,
@@ -342,7 +342,7 @@ class GeneratedNodeSchema(GeneratedBaseNodeSchema):
     )
 
 
-class GeneratedGenericSchema(GeneratedBaseNodeSchema):
+class GenericSchemaWrite(BaseNodeSchemaWrite):
     model_config = ConfigDict(extra="forbid")
     hierarchical: bool = Field(
         default=False,
@@ -356,3 +356,9 @@ class GeneratedGenericSchema(GeneratedBaseNodeSchema):
         default=None,
         description="Nodes inheriting from this Generic schema must belong to one of the listed namespaces",
     )
+
+
+class InfrahubSchemaWrite(BaseModel):
+    version: str | None = None
+    nodes: list[NodeSchemaWrite] = Field(default_factory=list)
+    generics: list[GenericSchemaWrite] = Field(default_factory=list)

infrahub_sdk/schema/validate.py

@@ -16,16 +16,16 @@
 from pydantic import ValidationError as PydanticValidationError
 
 from .generated.write import (
-    GeneratedAttributeSchema,
-    GeneratedGenericSchema,
-    GeneratedNodeSchema,
-    GeneratedRelationshipSchema,
+    AttributeSchemaWrite,
+    GenericSchemaWrite,
+    NodeSchemaWrite,
+    RelationshipSchemaWrite,
 )
 
 # Maps each collection in a schema-root payload to the write model its items must satisfy.
 _WRITE_MODELS_BY_COLLECTION: dict[str, type[BaseModel]] = {
-    "nodes": GeneratedNodeSchema,
-    "generics": GeneratedGenericSchema,
+    "nodes": NodeSchemaWrite,
+    "generics": GenericSchemaWrite,
 }
 
 
@@ -109,7 +109,7 @@ def _validate_extensions(extensions: Any, errors: list[SchemaValidationErrorDeta
         if isinstance(attributes, list):
             for attr_index, attribute in enumerate(attributes):
                 _validate_item(
-                    model=GeneratedAttributeSchema,
+                    model=AttributeSchemaWrite,
                     item=attribute,
                     prefix=f"{node_prefix}.attributes[{attr_index}]",
                     errors=errors,
@@ -118,7 +118,7 @@ def _validate_extensions(extensions: Any, errors: list[SchemaValidationErrorDeta
         if isinstance(relationships, list):
             for rel_index, relationship in enumerate(relationships):
                 _validate_item(
-                    model=GeneratedRelationshipSchema,
+                    model=RelationshipSchemaWrite,
                     item=relationship,
                     prefix=f"{node_prefix}.relationships[{rel_index}]",
                     errors=errors,

tests/unit/test_schema_generated_models.py

@@ -15,20 +15,23 @@
 
 import pytest
 
+from infrahub_sdk.schema import InfrahubSchemaRead, InfrahubSchemaWrite
 from infrahub_sdk.schema.generated import read as read_module
 from infrahub_sdk.schema.generated import write as write_module
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
 
 _GENERATED_DIR = Path(write_module.__file__).parent
-_EXPECTED_FAMILIES = [
-    "GeneratedAttributeSchema",
-    "GeneratedRelationshipSchema",
-    "GeneratedBaseNodeSchema",
-    "GeneratedNodeSchema",
-    "GeneratedGenericSchema",
+# Each family pairs its write-variant class name with its read-variant class name.
+_FAMILY_PAIRS = [
+    ("AttributeSchemaWrite", "AttributeSchemaRead"),
+    ("RelationshipSchemaWrite", "RelationshipSchemaRead"),
+    ("BaseNodeSchemaWrite", "BaseNodeSchemaRead"),
+    ("NodeSchemaWrite", "NodeSchemaRead"),
+    ("GenericSchemaWrite", "GenericSchemaRead"),
 ]
+_WRITE_FAMILIES = [write for write, _ in _FAMILY_PAIRS]
 
 
 @pytest.mark.parametrize("filename", ["write.py", "read.py"])
@@ -38,31 +41,39 @@ def test_generated_files_present_with_do_not_edit_header(filename: str) -> None:
     assert "do not edit" in path.read_text().splitlines()[0].lower()
 
 
-@pytest.mark.parametrize("family", _EXPECTED_FAMILIES)
-def test_expected_model_families_present_in_both_variants(family: str) -> None:
-    assert hasattr(write_module, family), f"write variant is missing {family}"
-    assert hasattr(read_module, family), f"read variant is missing {family}"
+@pytest.mark.parametrize(("write_family", "read_family"), _FAMILY_PAIRS)
+def test_expected_model_families_present_in_each_variant(write_family: str, read_family: str) -> None:
+    assert hasattr(write_module, write_family), f"write variant is missing {write_family}"
+    assert hasattr(read_module, read_family), f"read variant is missing {read_family}"
 
 
-@pytest.mark.parametrize("family", _EXPECTED_FAMILIES)
+@pytest.mark.parametrize("family", _WRITE_FAMILIES)
 def test_write_variant_forbids_extra_fields(family: str) -> None:
     model: type[BaseModel] = getattr(write_module, family)
     assert model.model_config.get("extra") == "forbid", (
         f"write variant {family} must set extra='forbid' so non-settable fields are rejected"
     )
 
 
-@pytest.mark.parametrize("family", _EXPECTED_FAMILIES)
-def test_read_variant_is_superset_of_write_variant(family: str) -> None:
-    write_model: type[BaseModel] = getattr(write_module, family)
-    read_model: type[BaseModel] = getattr(read_module, family)
+@pytest.mark.parametrize(("write_family", "read_family"), _FAMILY_PAIRS)
+def test_read_variant_is_superset_of_write_variant(write_family: str, read_family: str) -> None:
+    write_model: type[BaseModel] = getattr(write_module, write_family)
+    read_model: type[BaseModel] = getattr(read_module, read_family)
     write_fields = set(write_model.model_fields)
     read_fields = set(read_model.model_fields)
     missing = write_fields - read_fields
-    assert not missing, f"read variant {family} must expose every write field; missing: {sorted(missing)}"
+    assert not missing, f"read variant {read_family} must expose every write field; missing: {sorted(missing)}"
 
 
 def test_write_variant_carries_read_level_fields_absent() -> None:
     # `inherited` is a read-level attribute field and must not exist on the write variant.
-    assert "inherited" not in write_module.GeneratedAttributeSchema.model_fields
-    assert "inherited" in read_module.GeneratedAttributeSchema.model_fields
+    assert "inherited" not in write_module.AttributeSchemaWrite.model_fields
+    assert "inherited" in read_module.AttributeSchemaRead.model_fields
+
+
+def test_document_root_models_import_with_nodes_and_generics() -> None:
+    for root in (InfrahubSchemaWrite, InfrahubSchemaRead):
+        assert "nodes" in root.model_fields, f"{root.__name__} must expose a 'nodes' field"
+        assert "generics" in root.model_fields, f"{root.__name__} must expose a 'generics' field"
+    # The write root is not extra=forbid so schema-level keys outside the model are tolerated.
+    assert InfrahubSchemaWrite.model_config.get("extra") != "forbid"

tests/unit/test_schema_offline_validation.py

@@ -8,7 +8,7 @@
 
 import pytest
 
-from infrahub_sdk.schema import validate_schema
+from infrahub_sdk.schema import InfrahubSchemaRead, InfrahubSchemaWrite, validate_schema
 from infrahub_sdk.schema.validate import SchemaValidationResult
 
 
@@ -34,6 +34,12 @@ def _valid_schema() -> dict:
     }
 
 
+def test_schema_root_models_are_importable_with_nodes_and_generics() -> None:
+    for root in (InfrahubSchemaWrite, InfrahubSchemaRead):
+        assert "nodes" in root.model_fields
+        assert "generics" in root.model_fields
+
+
 def test_valid_payload_passes() -> None:
     result = validate_schema(schema=_valid_schema())
     assert isinstance(result, SchemaValidationResult)