fix: duck-type get_kind() instead of isinstance against write-side schema classes

infrahub_sync/generator/__init__.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Protocol
 
 import jinja2
 from infrahub_sdk.schema import (
-    AttributeSchema,
     NodeSchema,
     RelationshipKind,
-    RelationshipSchema,
 )
 
 if TYPE_CHECKING:
@@ -117,34 +115,54 @@ def get_children(node: NodeSchema, config: SyncConfig) -> str | None:
     return "{" + ", ".join(children_list) + "}"
 
 
-def get_kind(item: Union[RelationshipSchema, AttributeSchema]) -> str:
-    kind = "str"
-    if isinstance(item, AttributeSchema):
-        kind = ATTRIBUTE_KIND_MAP.get(item.kind, "str")
-        if item.optional:
-            kind = f"{kind} | None"
-            if item.default_value is not None:
-                # Format the default value based on its type
-                if isinstance(item.default_value, str):
-                    kind += f' = "{item.default_value}"'
-                elif isinstance(item.default_value, (int, float, bool)):
-                    kind += f" = {item.default_value}"
-                else:
-                    kind += f" = {item.default_value!r}"
+class _AttributeLike(Protocol):
+    """Structural shape get_attribute_type_annotation() needs from an attribute-schema object."""
+
+    kind: Any
+    optional: bool
+    default_value: Any
+
+
+class _RelationshipLike(Protocol):
+    """Structural shape get_relationship_type_annotation() needs from a relationship-schema object."""
+
+    cardinality: str
+    optional: bool
+
+
+def get_attribute_type_annotation(item: _AttributeLike) -> str:
+    """Return type annotation of schema attribute for Diffsync model."""
+    annotation = ATTRIBUTE_KIND_MAP.get(item.kind, "str")
+    if item.optional:
+        annotation = f"{annotation} | None"
+        if item.default_value is not None:
+            # Format the default value based on its type
+            if isinstance(item.default_value, str):
+                annotation += f' = "{item.default_value}"'
+            elif isinstance(item.default_value, (int, float, bool)):
+                annotation += f" = {item.default_value}"
             else:
-                kind += " = None"
+                annotation += f" = {item.default_value!r}"
+        else:
+            annotation += " = None"
+
+    return annotation
+
 
-    elif isinstance(item, RelationshipSchema) and item.cardinality == "one":
+def get_relationship_type_annotation(item: _RelationshipLike) -> str:
+    """Return type annotation of schema relationship for Diffsync model."""
+    annotation = "str"
+    if item.cardinality == "one":
         if item.optional:
-            kind = f"{kind} | None = None"
+            annotation = f"{annotation} | None = None"
 
-    elif isinstance(item, RelationshipSchema) and item.cardinality == "many":
-        kind = "list[str]"
+    elif item.cardinality == "many":
+        annotation = "list[str]"
         if item.optional:
-            kind = f"{kind} | None"
-        kind += " = []"
+            annotation = f"{annotation} | None"
+        annotation += " = []"
 
-    return kind
+    return annotation
 
 
 def has_children(node: NodeSchema, config: SyncConfig) -> bool:
@@ -165,7 +183,8 @@ def render_template(template_file: Path, output_dir: Path, output_file: Path, co
     template_env.filters["has_node"] = has_node
     template_env.filters["has_field"] = has_field
     template_env.filters["has_children"] = has_children
-    template_env.filters["get_kind"] = get_kind
+    template_env.filters["get_attribute_type_annotation"] = get_attribute_type_annotation
+    template_env.filters["get_relationship_type_annotation"] = get_relationship_type_annotation
 
     template = template_env.get_template(str(template_file))
 

infrahub_sync/generator/templates/diffsync_models.j2

@@ -35,12 +35,12 @@ class {{ nodekind }}(_ModelBaseClass):
 
 {%-      for attr in node.attributes -%}
 {%-          if config | has_field(node.kind, attr.name) %}
-    {{ attr.name }}: {{ attr | get_kind }}
+    {{ attr.name }}: {{ attr | get_attribute_type_annotation }}
 {%-          endif -%}
 {%-      endfor -%}
 {%-      for rel in node.relationships -%}
 {%-          if config | has_field(node.kind, rel.name) %}
-    {{ rel.name }}: {{ rel | get_kind }}
+    {{ rel.name }}: {{ rel | get_relationship_type_annotation }}
 {%-          endif -%}
 {%-      endfor %}
 

tests/test_generator.py

@@ -0,0 +1,96 @@
+"""Tests for infrahub_sync.generator.get_attribute_type_annotation / get_relationship_type_annotation.
+
+These functions decide the pydantic type annotation for each generated DiffSync
+model field. They must work on whatever object `client.schema.all()` hands back
+-- infrahub-sdk's read-side schema classes (AttributeSchemaAPI /
+RelationshipSchemaAPI) -- not just on the write-side AttributeSchema /
+RelationshipSchema classes used to build schema payloads.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+from infrahub_sdk.schema import AttributeSchema, RelationshipSchema
+from infrahub_sdk.schema.main import AttributeKind, AttributeSchemaAPI, RelationshipSchemaAPI
+
+from infrahub_sync.generator import get_attribute_type_annotation, get_relationship_type_annotation
+
+
+@dataclass
+class _FakeAttribute:
+    """Duck-types an attribute schema without subclassing any SDK class."""
+
+    kind: str
+    optional: bool = False
+    default_value: Any = None
+
+
+@dataclass
+class _FakeRelationship:
+    """Duck-types a relationship schema without subclassing any SDK class."""
+
+    cardinality: str
+    optional: bool = False
+
+
+def test_fake_schema_objects_do_not_subclass_sdk_schema_classes() -> None:
+    """Pins down the premise the other tests rely on.
+
+    If this starts failing, the fakes below stopped simulating the
+    decoupled-class scenario and no longer guard the regression.
+    """
+    assert not isinstance(_FakeAttribute(kind="Text"), AttributeSchema)
+    assert not isinstance(_FakeRelationship(cardinality="one"), RelationshipSchema)
+
+
+@pytest.mark.parametrize(
+    ("attribute", "expected"),
+    [
+        (_FakeAttribute(kind="Text"), "str"),
+        (_FakeAttribute(kind="Number"), "int"),
+        (_FakeAttribute(kind="Boolean"), "bool"),
+        (_FakeAttribute(kind="SomeUnmappedKind"), "str"),
+        (_FakeAttribute(kind="Text", optional=True), "str | None = None"),
+        (_FakeAttribute(kind="Text", optional=True, default_value="foo"), 'str | None = "foo"'),
+        (_FakeAttribute(kind="Number", optional=True, default_value=5), "int | None = 5"),
+        (_FakeAttribute(kind="Boolean", optional=True, default_value=True), "bool | None = True"),
+    ],
+)
+def test_get_attribute_type_annotation_for_attribute_shaped_objects(attribute: _FakeAttribute, expected: str) -> None:
+    assert get_attribute_type_annotation(attribute) == expected
+
+
+@pytest.mark.parametrize(
+    ("relationship", "expected"),
+    [
+        (_FakeRelationship(cardinality="one"), "str"),
+        (_FakeRelationship(cardinality="one", optional=True), "str | None = None"),
+        (_FakeRelationship(cardinality="many"), "list[str] = []"),
+        (_FakeRelationship(cardinality="many", optional=True), "list[str] | None = []"),
+    ],
+)
+def test_get_relationship_type_annotation_for_relationship_shaped_objects(
+    relationship: _FakeRelationship, expected: str
+) -> None:
+    assert get_relationship_type_annotation(relationship) == expected
+
+
+def test_get_attribute_type_annotation_against_real_sdk_read_side_schema_classes() -> None:
+    """Sanity check against the actual classes `client.schema.all()` returns."""
+    optional_attr = AttributeSchemaAPI(id="1", name="description", kind=AttributeKind.TEXT, optional=True)
+    assert get_attribute_type_annotation(optional_attr) == "str | None = None"
+
+    required_attr = AttributeSchemaAPI(id="2", name="name", kind=AttributeKind.TEXT, optional=False)
+    assert get_attribute_type_annotation(required_attr) == "str"
+
+
+def test_get_relationship_type_annotation_against_real_sdk_read_side_schema_classes() -> None:
+    """Sanity check against the actual classes `client.schema.all()` returns."""
+    many_rel = RelationshipSchemaAPI(id="3", name="tags", peer="BuiltinTag", cardinality="many", optional=True)
+    assert get_relationship_type_annotation(many_rel) == "list[str] | None = []"
+
+    one_rel = RelationshipSchemaAPI(id="4", name="status", peer="StatusGeneric", cardinality="one", optional=False)
+    assert get_relationship_type_annotation(one_rel) == "str"