diff --git a/docs/changes/newsfragments/502.bugfix b/docs/changes/newsfragments/502.bugfix new file mode 100644 index 000000000..f88e949df --- /dev/null +++ b/docs/changes/newsfragments/502.bugfix @@ -0,0 +1 @@ +Fix type issues and update model configs for :class:`.BaseDataGrabber`, :class:`.BasePreprocessor`, :class:`.BaseMarker` and :class:`.BaseStorage` by `Synchon Mandal`_ diff --git a/junifer/datagrabber/base.py b/junifer/datagrabber/base.py index c34fabdb5..4fb3796a6 100644 --- a/junifer/datagrabber/base.py +++ b/junifer/datagrabber/base.py @@ -13,7 +13,7 @@ from typing import Annotated, Any import structlog from aenum import Enum as AEnum -from pydantic import BaseModel, BeforeValidator, ConfigDict, Field +from pydantic import BaseModel, BeforeValidator, ConfigDict from ..pipeline import UpdateMetaMixin 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[ DataType | list[DataType], - Field(frozen=True), BeforeValidator(ensure_list), ] datadir: Path diff --git a/junifer/markers/base.py b/junifer/markers/base.py index be2ecba1c..06cb48fd8 100644 --- a/junifer/markers/base.py +++ b/junifer/markers/base.py @@ -6,16 +6,16 @@ from abc import ABC, abstractmethod from copy import deepcopy -from typing import Any, ClassVar +from typing import Annotated, Any, ClassVar import structlog -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, BeforeValidator, ConfigDict from ..datagrabber import DataType from ..pipeline import PipelineStepMixin, UpdateMetaMixin from ..storage import StorageType from ..typing import MarkerInOutMappings, StorageLike -from ..utils import raise_error +from ..utils import ensure_list_or_none, raise_error __all__ = ["BaseMarker"] @@ -55,9 +55,12 @@ class BaseMarker(BaseModel, ABC, PipelineStepMixin, UpdateMetaMixin): _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 def model_post_init(self, context: Any): # noqa: D102 diff --git a/junifer/preprocess/base.py b/junifer/preprocess/base.py index fd84c38a5..6c3a41e41 100644 --- a/junifer/preprocess/base.py +++ b/junifer/preprocess/base.py @@ -55,7 +55,7 @@ class BasePreprocessor(BaseModel, ABC, PipelineStepMixin, UpdateMetaMixin): _VALID_DATA_TYPES: ClassVar[Sequence[DataType]] - model_config = ConfigDict(use_enum_values=True) + model_config = ConfigDict(extra="forbid", use_enum_values=True) on: Annotated[ DataType | list[DataType] | None, diff --git a/junifer/storage/base.py b/junifer/storage/base.py index 2408cd0bf..a5c862db1 100644 --- a/junifer/storage/base.py +++ b/junifer/storage/base.py @@ -65,7 +65,9 @@ class BaseFeatureStorage(BaseModel, ABC): _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 single_output: bool = True