Merge pull request #23 from NowanIlfideme/fix/mypy-utils

Fixes for mypy in utils
NowanIlfideme · May 10, 2023 · f35316d · f35316d
2 parents e6a1173 + 2300b16
commit f35316d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/pydantic_kedro/utils.py b/src/pydantic_kedro/utils.py
@@ -2,6 +2,7 @@
 
 from typing import Literal, Type, TypeVar
 
+from kedro.io.core import AbstractDataSet
 from pydantic import BaseModel
 
 from pydantic_kedro.datasets.auto import PydanticAutoDataSet
@@ -15,7 +16,7 @@
 T = TypeVar("T", bound=BaseModel)
 
 
-def load_model(uri: str, supercls: Type[T] = BaseModel) -> T:
+def load_model(uri: str, supercls: Type[T] = BaseModel) -> T:  # type: ignore
     """Load a Pydantic model from a given URI.
 
     Parameters
@@ -52,6 +53,7 @@ def save_model(
     """
     if not isinstance(model, BaseModel):
         raise TypeError(f"Expected Pydantic model, but got {model!r}")
+    ds: AbstractDataSet
     if format == "auto":
         ds = PydanticAutoDataSet(uri)
     elif format == "zip":