Refactor parser base post-processing for DRY and type-safe implementation (#2730)

Author: koxudaxi

src/datamodel_code_generator/imports.py

@@ -168,6 +168,33 @@ def dump_all(self, *, multiline: bool = False) -> str:
         items = ", ".join(f'"{name}"' for name in name_list)
         return f"__all__ = [{items}]"
 
+    def get_effective_name(self, from_: str | None, import_: str) -> str:
+        """Get the effective name after alias resolution."""
+        return self.alias.get(from_, {}).get(import_, import_)
+
+    def remove_unused(self, used_names: set[str]) -> None:
+        """Remove imports not referenced in used_names.
+
+        Note: Checks both effective name (after alias) and original name to handle
+        cases where code may reference either form (e.g., type annotations may use
+        original name while runtime code uses alias).
+        """
+        unused = [
+            (from_, import_)
+            for from_, imports_ in self.items()
+            for import_ in imports_
+            if not {self.get_effective_name(from_, import_), import_}.intersection(used_names)
+        ]
+        for from_, import_ in unused:
+            alias = self.alias.get(from_, {}).get(import_)
+            reference_path = next(
+                (p for p, i in self.reference_paths.items() if i.from_ == from_ and i.import_ == import_),
+                None,
+            )
+            import_obj = Import(from_=from_, import_=import_, alias=alias, reference_path=reference_path)
+            while self.counter.get((from_, import_), 0) > 0:
+                self.remove(import_obj)
+
 
 IMPORT_ANNOTATED = Import.from_full_path("typing.Annotated")
 IMPORT_ANY = Import.from_full_path("typing.Any")

src/datamodel_code_generator/model/base.py

@@ -478,6 +478,10 @@ class DataModel(TemplateBase, Nullable, ABC):  # noqa: PLR0904
     DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = ()
     IS_ALIAS: ClassVar[bool] = False
     SUPPORTS_GENERIC_BASE_CLASS: ClassVar[bool] = True
+    SUPPORTS_DISCRIMINATOR: ClassVar[bool] = False
+    SUPPORTS_FIELD_RENAMING: ClassVar[bool] = False
+    SUPPORTS_WRAPPED_DEFAULT: ClassVar[bool] = False
+    SUPPORTS_KW_ONLY: ClassVar[bool] = False
     has_forward_reference: bool = False
 
     def __init__(  # noqa: PLR0913

src/datamodel_code_generator/model/dataclass.py

@@ -42,6 +42,8 @@ class DataClass(DataModel):
 
     TEMPLATE_FILE_PATH: ClassVar[str] = "dataclass.jinja2"
     DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = (IMPORT_DATACLASS,)
+    SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
+    SUPPORTS_KW_ONLY: ClassVar[bool] = True
 
     def __init__(  # noqa: PLR0913
         self,

src/datamodel_code_generator/model/msgspec.py

@@ -113,6 +113,7 @@ class Struct(DataModel):
     BASE_CLASS_NAME: ClassVar[str] = "Struct"
     BASE_CLASS_ALIAS: ClassVar[str] = "_Struct"
     DEFAULT_IMPORTS: ClassVar[tuple[Import, ...]] = ()
+    SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
     CONFIG_MAPPING: ClassVar[dict[tuple[str, Any], tuple[str, Any] | None]] = {
         ("allow_mutation", False): ("frozen", True),
         ("extra_fields", "forbid"): ("forbid_unknown_fields", True),

src/datamodel_code_generator/model/pydantic/base_model.py

@@ -328,6 +328,7 @@ class BaseModel(BaseModelBase):
 
     TEMPLATE_FILE_PATH: ClassVar[str] = "pydantic/BaseModel.jinja2"
     BASE_CLASS: ClassVar[str] = "pydantic.BaseModel"
+    SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
 
     def __init__(  # noqa: PLR0912, PLR0913
         self,

src/datamodel_code_generator/model/pydantic_v2/base_model.py

@@ -174,6 +174,9 @@ class BaseModel(BaseModelBase):
     BASE_CLASS: ClassVar[str] = "pydantic.BaseModel"
     BASE_CLASS_NAME: ClassVar[str] = "BaseModel"
     BASE_CLASS_ALIAS: ClassVar[str] = "_BaseModel"
+    SUPPORTS_DISCRIMINATOR: ClassVar[bool] = True
+    SUPPORTS_FIELD_RENAMING: ClassVar[bool] = True
+    SUPPORTS_WRAPPED_DEFAULT: ClassVar[bool] = True
     CONFIG_ATTRIBUTES: ClassVar[list[ConfigAttribute]] = [
         ConfigAttribute("allow_population_by_field_name", "populate_by_name", False),  # noqa: FBT003
         ConfigAttribute("populate_by_name", "populate_by_name", False),  # noqa: FBT003