diff --git a/infrahub_sync/generator/__init__.py b/infrahub_sync/generator/__init__.py index be82a795..1cf6c9d6 100644 --- a/infrahub_sync/generator/__init__.py +++ b/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)) diff --git a/infrahub_sync/generator/templates/diffsync_models.j2 b/infrahub_sync/generator/templates/diffsync_models.j2 index 577bff86..37321ff5 100644 --- a/infrahub_sync/generator/templates/diffsync_models.j2 +++ b/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 %} diff --git a/tests/test_generator.py b/tests/test_generator.py new file mode 100644 index 00000000..0c2267e9 --- /dev/null +++ b/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"