From 6475d4da9fc7a4dab0bf27895fae3002c8519eb5 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Wed, 8 Jul 2026 02:12:09 +0900 Subject: [PATCH 1/2] Type reference sources --- src/datamodel_code_generator/reference.py | 30 +++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/src/datamodel_code_generator/reference.py b/src/datamodel_code_generator/reference.py index 0dae4fc93..5c524f99e 100644 --- a/src/datamodel_code_generator/reference.py +++ b/src/datamodel_code_generator/reference.py @@ -43,7 +43,7 @@ import inflect - from datamodel_code_generator.model.base import DataModel + from datamodel_code_generator.model.base import DataModel, DataModelFieldBase from datamodel_code_generator.types import DataType @@ -77,6 +77,32 @@ def reference(self) -> Reference | None: ... +if TYPE_CHECKING: + + class _ReferenceSource(Protocol): + """Protocol for objects that can be assigned to Reference.source.""" + + fields: list[DataModelFieldBase] + + @property + def is_alias(self) -> bool: + """Return whether this source renders as an alias.""" + ... + + @property + def module_name(self) -> str | None: + """Return this source module name.""" + ... + + @property + def nullable(self) -> bool: + """Return whether this source is nullable.""" + ... + +else: + _ReferenceSource = ReferenceChild + + class _BaseModel(BaseModel): """Base model with field exclusion and pass-through support.""" @@ -124,7 +150,7 @@ class Reference(_BaseModel): name: str duplicate_name: Optional[str] = None # noqa: UP045 loaded: bool = True - source: Optional[ReferenceChild] = None # noqa: UP045 + source: Optional[_ReferenceSource] = None # noqa: UP045 children: list[ReferenceChild] = Field(default_factory=list) _exclude_fields: ClassVar[set[str]] = {"children"} From 10ce1c8b7e96b36e8f4c0dc79bca5aa7ddb0f136 Mon Sep 17 00:00:00 2001 From: Koudai Aono Date: Wed, 8 Jul 2026 02:31:59 +0900 Subject: [PATCH 2/2] Fix reference source typing --- src/datamodel_code_generator/reference.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/datamodel_code_generator/reference.py b/src/datamodel_code_generator/reference.py index 5c524f99e..b558f6ada 100644 --- a/src/datamodel_code_generator/reference.py +++ b/src/datamodel_code_generator/reference.py @@ -43,7 +43,6 @@ import inflect - from datamodel_code_generator.model.base import DataModel, DataModelFieldBase from datamodel_code_generator.types import DataType @@ -54,7 +53,7 @@ def _is_data_type(value: object) -> TypeIs[DataType]: return isinstance(value, DataType_) -def _is_data_model(value: object) -> TypeIs[DataModel]: +def _is_data_model(value: object) -> bool: """Check if value is a DataModel instance.""" from datamodel_code_generator.model.base import DataModel as DataModel_ # noqa: PLC0415 @@ -79,25 +78,25 @@ def reference(self) -> Reference | None: if TYPE_CHECKING: - class _ReferenceSource(Protocol): + class _ReferenceSource(ReferenceChild, Protocol): """Protocol for objects that can be assigned to Reference.source.""" - fields: list[DataModelFieldBase] + fields: Sequence[Any] @property def is_alias(self) -> bool: """Return whether this source renders as an alias.""" - ... + raise NotImplementedError @property def module_name(self) -> str | None: """Return this source module name.""" - ... + raise NotImplementedError @property def nullable(self) -> bool: """Return whether this source is nullable.""" - ... + raise NotImplementedError else: _ReferenceSource = ReferenceChild @@ -184,7 +183,7 @@ def replace_children_references(self, new_reference: Reference) -> None: if _is_data_type(child): child.replace_reference(new_reference) - def iter_data_model_children(self) -> Iterator[DataModel]: + def iter_data_model_children(self) -> Iterator[Any]: """Yield all DataModel children.""" for child in self.children: if _is_data_model(child):