[BUG]: Fix type issues for base models #502

Merged
synchon merged 4 commits from fix/dg-base-types-frozen into main 2026-07-17 12:23:21 +00:00
5 changed files with 15 additions and 10 deletions

View file

@ -0,0 +1 @@
Fix type issues and update model configs for :class:`.BaseDataGrabber`, :class:`.BasePreprocessor`, :class:`.BaseMarker` and :class:`.BaseStorage` by `Synchon Mandal`_

View file

@ -13,7 +13,7 @@ from typing import Annotated, Any
import structlog import structlog
from aenum import Enum as AEnum from aenum import Enum as AEnum
from pydantic import BaseModel, BeforeValidator, ConfigDict, Field from pydantic import BaseModel, BeforeValidator, ConfigDict
from ..pipeline import UpdateMetaMixin from ..pipeline import UpdateMetaMixin
from ..typing import Element, Elements from ..typing import Element, Elements
@ -58,11 +58,10 @@ class BaseDataGrabber(BaseModel, ABC, UpdateMetaMixin):
""" """
model_config = ConfigDict(use_enum_values=True) model_config = ConfigDict(extra="forbid", use_enum_values=True)
types: Annotated[ types: Annotated[
DataType | list[DataType], DataType | list[DataType],
Field(frozen=True),
BeforeValidator(ensure_list), BeforeValidator(ensure_list),
] ]
datadir: Path datadir: Path

View file

@ -6,16 +6,16 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from copy import deepcopy from copy import deepcopy
from typing import Any, ClassVar from typing import Annotated, Any, ClassVar
import structlog import structlog
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, BeforeValidator, ConfigDict
from ..datagrabber import DataType from ..datagrabber import DataType
from ..pipeline import PipelineStepMixin, UpdateMetaMixin from ..pipeline import PipelineStepMixin, UpdateMetaMixin
from ..storage import StorageType from ..storage import StorageType
from ..typing import MarkerInOutMappings, StorageLike from ..typing import MarkerInOutMappings, StorageLike
from ..utils import raise_error from ..utils import ensure_list_or_none, raise_error
__all__ = ["BaseMarker"] __all__ = ["BaseMarker"]
@ -55,9 +55,12 @@ class BaseMarker(BaseModel, ABC, PipelineStepMixin, UpdateMetaMixin):
_MARKER_INOUT_MAPPINGS: ClassVar[MarkerInOutMappings] _MARKER_INOUT_MAPPINGS: ClassVar[MarkerInOutMappings]
model_config = ConfigDict(use_enum_values=True) model_config = ConfigDict(extra="forbid", use_enum_values=True)
on: list[DataType] | None = None on: Annotated[
DataType | list[DataType] | None,
BeforeValidator(ensure_list_or_none),
] = None
name: str | None = None name: str | None = None
def model_post_init(self, context: Any): # noqa: D102 def model_post_init(self, context: Any): # noqa: D102

View file

@ -55,7 +55,7 @@ class BasePreprocessor(BaseModel, ABC, PipelineStepMixin, UpdateMetaMixin):
_VALID_DATA_TYPES: ClassVar[Sequence[DataType]] _VALID_DATA_TYPES: ClassVar[Sequence[DataType]]
model_config = ConfigDict(use_enum_values=True) model_config = ConfigDict(extra="forbid", use_enum_values=True)
on: Annotated[ on: Annotated[
DataType | list[DataType] | None, DataType | list[DataType] | None,

View file

@ -65,7 +65,9 @@ class BaseFeatureStorage(BaseModel, ABC):
_STORAGE_TYPES: ClassVar[Sequence[StorageType]] _STORAGE_TYPES: ClassVar[Sequence[StorageType]]
model_config = ConfigDict(frozen=True, use_enum_values=True) model_config = ConfigDict(
extra="forbid", frozen=True, use_enum_values=True
)
uri: Path uri: Path
single_output: bool = True single_output: bool = True