fix: types

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

src/bentoml/_internal/io_descriptors/base.py

@@ -38,7 +38,7 @@
 
 
 @singledispatchmethod
-def set_sample(self: IODescriptor[t.Any], value: t.Any) -> IODescriptor[t.Any]:
+def set_sample(self: IODescriptor[t.Any], value: t.Any) -> None:
     raise InvalidArgument(
         f"Unsupported sample type: '{type(value)}' (value: {value}). To register type '{type(value)}' to {self.__class__.__name__} implement a dispatch function and register types to 'set_sample.register'"
     )
@@ -79,14 +79,7 @@ class IODescriptor(ABC, _OpenAPIMeta, t.Generic[IOType]):
     endpoint IO descriptor types in BentoServer.
     """
 
-    __slots__ = (
-        "_initialized",
-        "_args",
-        "_kwargs",
-        "_proto_fields",
-        "_mime_type",
-        "descriptor_id",
-    )
+    __slots__ = ("_args", "_kwargs", "_proto_fields", "_mime_type", "descriptor_id")
 
     HTTP_METHODS = ["POST"]
 
@@ -96,7 +89,7 @@ class IODescriptor(ABC, _OpenAPIMeta, t.Generic[IOType]):
     _rpc_content_type: str = "application/grpc"
     _proto_fields: tuple[ProtoField]
     _sample: IOType | None = None
-    _initialized: bool
+    _initialized: bool = False
     _args: t.Sequence[t.Any]
     _kwargs: dict[str, t.Any]
 
@@ -108,23 +101,22 @@ def __init_subclass__(cls, *, descriptor_id: str | None = None):
                 )
             IO_DESCRIPTOR_REGISTRY[descriptor_id] = cls
         cls.descriptor_id = descriptor_id
-        cls._initialized = False
 
     def __new__(cls, *args: t.Any, **kwargs: t.Any) -> Self:
         sample = kwargs.pop("_sample", None)
         klass = object.__new__(cls)
         if sample is None:
             set_sample.register(type(None), lambda self, _: self)
-        kls = klass._set_sample(sample)
-        kls._args = args
-        kls._kwargs = kwargs
-        return kls
+        klass._set_sample(sample)
+        klass._args = args
+        klass._kwargs = kwargs
+        return klass
 
     def __getattr__(self, name: str) -> t.Any:
         if not self._initialized:
             self._lazy_init()
         assert self._initialized
-        return getattr(self, name)
+        return object.__getattribute__(self, name)
 
     def __repr__(self) -> str:
         return self.__class__.__qualname__
@@ -135,7 +127,7 @@ def _lazy_init(self) -> None:
         del self._args
         del self._kwargs
 
-    _set_sample: singledispatchmethod[IODescriptor[t.Any]] = set_sample
+    _set_sample: singledispatchmethod[None] = set_sample
 
     @property
     def sample(self) -> IOType | None:

src/bentoml/_internal/io_descriptors/file.py

@@ -143,13 +143,11 @@ def from_sample(cls, sample: FileType | str, kind: FileKind = "binaryio") -> Sel
     @set_sample.register(type(FileLike))
     def _(cls, sample: FileLike[bytes]):
         cls.sample = sample
-        return cls
 
     @set_sample.register(t.IO)
     def _(cls, sample: t.IO[t.Any]):
         if isinstance(cls, File):
             cls.sample = FileLike[bytes](sample, "<sample>")
-        return cls
 
     @set_sample.register(str)
     @set_sample.register(os.PathLike)
@@ -160,7 +158,6 @@ def _(cls, sample: str):
             p = resolve_user_filepath(sample, ctx=None)
             with open(p, "rb") as f:
                 cls.sample = FileLike[bytes](f, "<sample>")
-        return cls
 
     @classmethod
     def from_spec(cls, spec: dict[str, t.Any]) -> Self:

src/bentoml/_internal/io_descriptors/image.py

@@ -240,15 +240,13 @@ def from_sample(
             def _(cls: Self, sample: ext.NpNDArray):
                 if isinstance(cls, Image):
                     cls.sample = PIL.Image.fromarray(sample, mode=pilmode)
-                return cls
 
         elif LazyType["PIL.Image.Image"]("PIL.Image.Image").isinstance(sample):
 
             @set_sample.register(PIL.Image.Image)
             def _(cls: Self, sample: PIL.Image.Image):
                 if isinstance(cls, Image):
                     cls.sample = sample
-                return cls
 
         return super().from_sample(
             sample,
@@ -266,7 +264,6 @@ def _(cls, sample: str):
                     cls.sample = PIL.Image.open(f)
             except PIL.UnidentifiedImageError as err:
                 raise BadInput(f"Failed to parse sample image file: {err}") from None
-        return cls
 
     def to_spec(self) -> dict[str, t.Any]:
         return {

src/bentoml/_internal/io_descriptors/json.py

@@ -217,7 +217,6 @@ def from_sample(
             def _(cls: Self, sample: pydantic.BaseModel):
                 if isinstance(cls, JSON):
                     cls.sample = sample
-                return cls
 
         return super().from_sample(
             sample, pydantic_model=pydantic_model, json_encoder=json_encoder
@@ -227,13 +226,11 @@ def _(cls: Self, sample: pydantic.BaseModel):
     def _(cls, sample: dict[str, t.Any]):
         if isinstance(cls, JSON):
             cls.sample = sample
-        return cls
 
     @set_sample.register(str)
     def _(cls, sample: str):
         if isinstance(cls, JSON):
             cls.sample = json.loads(sample)
-        return cls
 
     def to_spec(self) -> dict[str, t.Any]:
         return {

src/bentoml/_internal/io_descriptors/numpy.py

@@ -442,7 +442,6 @@ def _(cls, sample: ext.NpNDArray):
             cls.sample = sample
             cls._shape = sample.shape
             cls._dtype = sample.dtype
-        return cls
 
     @set_sample.register(list)
     @set_sample.register(tuple)
@@ -452,7 +451,6 @@ def _(cls, sample: t.Sequence[t.Any]):
             cls.sample = __
             cls._shape = __.shape
             cls._dtype = __.dtype
-        return cls
 
     async def from_proto(self, field: pb.NDArray | bytes) -> ext.NpNDArray:
         """

src/bentoml/_internal/io_descriptors/pandas.py

@@ -408,8 +408,6 @@ def _(cls: Self, sample: ext.NpNDArray):
                     cls._shape = __.shape
                     cls._columns = [str(i) for i in range(sample.shape[1])]
 
-                return cls
-
         return super().from_sample(
             sample,
             dtype=True,  # set to True to infer from given input
@@ -426,7 +424,6 @@ def _(cls, sample: pd.DataFrame):
             cls.sample = sample
             cls._shape = sample.shape
             cls._columns = [str(x) for x in list(sample.columns)]
-        return cls
 
     @set_sample.register(str)
     @set_sample.register(os.PathLike)
@@ -464,7 +461,6 @@ def _(cls, sample: str):
                 raise InvalidArgument(
                     f"Failed to create a 'pd.DataFrame' from sample {sample}: {e}"
                 ) from None
-        return cls
 
     def _convert_dtype(
         self, value: ext.PdDTypeArg | None
@@ -926,7 +922,6 @@ def _(cls: Self, sample: ext.NpNDArray):
                     cls.sample = __
                     cls._dtype = __.dtype
                     cls._shape = __.shape
-                return cls
 
         return super().from_sample(
             sample,
@@ -941,7 +936,6 @@ def _(cls, sample: ext.PdSeries):
             cls.sample = sample
             cls._dtype = sample.dtype
             cls._shape = sample.shape
-        return cls
 
     @set_sample.register(list)
     @set_sample.register(tuple)
@@ -952,7 +946,6 @@ def _(cls, sample: t.Sequence[t.Any]):
             cls.sample = __
             cls._dtype = __.dtype
             cls._shape = __.shape
-        return cls
 
     def input_type(self) -> LazyType[ext.PdSeries]:
         return LazyType("pandas", "Series")

src/bentoml/_internal/io_descriptors/text.py

@@ -108,13 +108,11 @@ def from_sample(cls, sample: str | bytes) -> Self:
     def _(cls, sample: str):
         if isinstance(cls, Text):
             cls.sample = sample
-        return cls
 
     @set_sample.register(bytes)
     def _(cls, sample: bytes):
         if isinstance(cls, Text):
             cls.sample = sample.decode("utf-8")
-        return cls
 
     def input_type(self) -> t.Type[str]:
         return str

src/bentoml/_internal/service/openapi/specification.py

@@ -104,7 +104,7 @@ class Schema:
     anyOf: t.Optional[t.List[Schema]] = None
     not_: t.Optional[Schema] = None
     items: t.Optional[t.Union[Schema, t.List[Schema]]] = None
-    properties: t.Optional[t.Dict[str, Schema | Reference]] = None
+    properties: t.Optional[t.Dict[str, t.Union[Schema, Reference]]] = None
     additionalProperties: t.Optional[t.Union[Schema, Reference, bool]] = None
     description: t.Optional[str] = None
     format: t.Optional[str] = None

tests/unit/_internal/io/test_numpy.py

@@ -85,6 +85,8 @@ def test_numpy_openapi_request_body():
 
     nparray = NumpyNdarray(dtype="float")
     nparray.sample_input = ExampleGeneric("asdf")  # type: ignore (test exception)
+    with pytest.raises(BadInput):
+        nparray.openapi_example()
 
 
 def test_numpy_openapi_responses():