4343
4444 import inflect
4545
46- from datamodel_code_generator .model .base import DataModel
4746 from datamodel_code_generator .types import DataType
4847
4948
@@ -54,7 +53,7 @@ def _is_data_type(value: object) -> TypeIs[DataType]:
5453 return isinstance (value , DataType_ )
5554
5655
57- def _is_data_model (value : object ) -> TypeIs [ DataModel ] :
56+ def _is_data_model (value : object ) -> bool :
5857 """Check if value is a DataModel instance."""
5958 from datamodel_code_generator .model .base import DataModel as DataModel_ # noqa: PLC0415
6059
@@ -77,6 +76,32 @@ def reference(self) -> Reference | None:
7776 ...
7877
7978
79+ if TYPE_CHECKING :
80+
81+ class _ReferenceSource (ReferenceChild , Protocol ):
82+ """Protocol for objects that can be assigned to Reference.source."""
83+
84+ fields : Sequence [Any ]
85+
86+ @property
87+ def is_alias (self ) -> bool :
88+ """Return whether this source renders as an alias."""
89+ raise NotImplementedError
90+
91+ @property
92+ def module_name (self ) -> str | None :
93+ """Return this source module name."""
94+ raise NotImplementedError
95+
96+ @property
97+ def nullable (self ) -> bool :
98+ """Return whether this source is nullable."""
99+ raise NotImplementedError
100+
101+ else :
102+ _ReferenceSource = ReferenceChild
103+
104+
80105class _BaseModel (BaseModel ):
81106 """Base model with field exclusion and pass-through support."""
82107
@@ -124,7 +149,7 @@ class Reference(_BaseModel):
124149 name : str
125150 duplicate_name : Optional [str ] = None # noqa: UP045
126151 loaded : bool = True
127- source : Optional [ReferenceChild ] = None # noqa: UP045
152+ source : Optional [_ReferenceSource ] = None # noqa: UP045
128153 children : list [ReferenceChild ] = Field (default_factory = list )
129154 _exclude_fields : ClassVar [set [str ]] = {"children" }
130155
@@ -158,7 +183,7 @@ def replace_children_references(self, new_reference: Reference) -> None:
158183 if _is_data_type (child ):
159184 child .replace_reference (new_reference )
160185
161- def iter_data_model_children (self ) -> Iterator [DataModel ]:
186+ def iter_data_model_children (self ) -> Iterator [Any ]:
162187 """Yield all DataModel children."""
163188 for child in self .children :
164189 if _is_data_model (child ):
0 commit comments