dynamic types support array, dict, ndarray

mlem/contrib/numpy.py

@@ -94,7 +94,7 @@ class NumpyNdarrayType(
  type: ClassVar[str] = "ndarray"
  libraries: ClassVar[List[ModuleType]] = [np]
 
- shape: Tuple[Optional[int], ...]
+ shape: Optional[Tuple[Optional[int], ...]]
  dtype: str
 
  @staticmethod
@@ -134,7 +134,16 @@ def _subtype(self, subshape: Tuple[Optional[int], ...]):
  def get_model(self, prefix: str = "") -> Type[BaseModel]:
  # TODO: https://github.com/iterative/mlem/issues/33
  return create_model(
- prefix + "NumpyNdarray", __root__=(List[self._subtype(self.shape[1:])], ...) # type: ignore
+ prefix + "NumpyNdarray",
+ __root__=(
+ self._subtype(self.shape)
+ if self.shape is not None
+ else List[
+ Union[python_type_from_np_string_repr(self.dtype), List]
+ ], # type: ignore
+ ...,
+ )
+ # type: ignore
  )
 
  def serialize(self, instance: np.ndarray):
@@ -148,10 +157,20 @@ def serialize(self, instance: np.ndarray):
  return instance.tolist()
 
  def _check_shape(self, array, exc_type):
- if tuple(array.shape)[1:] != self.shape[1:]:
- raise exc_type(
- f"given array is of shape: {(None,) + tuple(array.shape)[1:]}, expected: {self.shape}"
+ if self.shape is not None:
+ if len(array.shape) != len(self.shape):
+ raise exc_type(
+ f"given array is of rank: {len(array.shape)}, expected: {len(self.shape)}"
+ )
+
+ array_shape = tuple(
+ None if expected_dim is None else array_dim
+ for array_dim, expected_dim in zip(array.shape, self.shape)
  )
+ if tuple(array_shape) != self.shape:
+ raise exc_type(
+ f"given array is of shape: {array_shape}, expected: {self.shape}"
+ )
 
  def get_writer(self, project: str = None, filename: str = None, **kwargs):
  return NumpyArrayWriter()

mlem/core/data_type.py

@@ -2,6 +2,7 @@
 Base classes for working with data in MLEM
 """
 import builtins
+import json
 import posixpath
 from abc import ABC, abstractmethod
 from typing import (
@@ -18,7 +19,7 @@
 )
 
 import flatdict
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
 from pydantic.main import create_model
 
 from mlem.core.artifacts import Artifacts, Storage
@@ -70,8 +71,10 @@ def bind(self, data: Any):
  return self
 
  @classmethod
- def create(cls, obj: Any, **kwargs):
- return DataAnalyzer.analyze(obj, **kwargs).bind(obj)
+ def create(cls, obj: Any, is_dynamic: bool = False, **kwargs):
+ return DataAnalyzer.analyze(obj, is_dynamic=is_dynamic, **kwargs).bind(
+ obj
+ )
 
 
 class DataSerializer(ABC):
@@ -367,7 +370,7 @@ def get_model(self, prefix: str = "") -> Type[BaseModel]:
 
 def _check_type_and_size(obj, dtype, size, exc_type):
  DataType.check_type(obj, dtype, exc_type)
- if size != -1 and len(obj) != size:
+ if size is not None and len(obj) != size:
  raise exc_type(
  f"given {dtype.__name__} has len: {len(obj)}, expected: {size}"
  )
@@ -449,44 +452,77 @@ def is_object_valid(cls, obj: Any) -> bool:
  return isinstance(obj, (list, tuple))
 
  @classmethod
- def process(cls, obj, **kwargs) -> DataType:
+ def process(cls, obj, is_dynamic: bool = False, **kwargs) -> DataType:
  if isinstance(obj, tuple):
- return TupleType(items=[DataAnalyzer.analyze(o) for o in obj])
+ return TupleType(
+ items=[
+ DataAnalyzer.analyze(o, is_dynamic=is_dynamic, **kwargs)
+ for o in obj
+ ]
+ )
 
  py_types = {type(o) for o in obj}
  if len(obj) <= 1 or len(py_types) > 1:
- return ListType(items=[DataAnalyzer.analyze(o) for o in obj])
+ return ListType(
+ items=[
+ DataAnalyzer.analyze(o, is_dynamic=is_dynamic, **kwargs)
+ for o in obj
+ ]
+ )
+
+ size = None if is_dynamic else len(obj)
 
  if not py_types.intersection(
  PrimitiveType.PRIMITIVES
  ): # py_types is guaranteed to be singleton set here
- items_types = [DataAnalyzer.analyze(o) for o in obj]
+ items_types = [
+ DataAnalyzer.analyze(o, is_dynamic=is_dynamic, **kwargs)
+ for o in obj
+ ]
  first, *others = items_types
  for other in others:
  if first != other:
  return ListType(items=items_types)
- return ArrayType(dtype=first, size=len(obj))
+ return ArrayType(dtype=first, size=size)
 
  # optimization for large lists of same primitive type elements
- return ArrayType(dtype=DataAnalyzer.analyze(obj[0]), size=len(obj))
+ return ArrayType(
+ dtype=DataAnalyzer.analyze(
+ obj[0], is_dynamic=is_dynamic, **kwargs
+ ),
+ size=size,
+ )
+
+
+class DictTypeHook(DataHook):
+ @classmethod
+ def is_object_valid(cls, obj: Any) -> bool:
+ return isinstance(obj, dict)
+
+ @classmethod
+ def process(
+ cls, obj: Any, is_dynamic: bool = False, **kwargs
+ ) -> Union["DictType", "DynamicDictType"]:
+ if not is_dynamic:
+ return DictType.create(obj, **kwargs)
+ return DynamicDictType.create(obj, **kwargs)
 
 
-class DictType(DataType, DataSerializer, DataHook):
+class DictType(DataType, DataSerializer):
  """
- DataType for dict
+ DataType for dict with fixed set of keys
  """
 
  type: ClassVar[str] = "dict"
  item_types: Dict[str, DataType]
 
  @classmethod
- def is_object_valid(cls, obj: Any) -> bool:
- return isinstance(obj, dict)
-
- @classmethod
- def process(cls, obj: Any, **kwargs) -> "DictType":
+ def create(cls, obj, **kwargs):
  return DictType(
- item_types={k: DataAnalyzer.analyze(v) for (k, v) in obj.items()}
+ item_types={
+ k: DataAnalyzer.analyze(v, is_dynamic=False, **kwargs)
+ for (k, v) in obj.items()
+ }
  )
 
  def deserialize(self, obj):
@@ -577,6 +613,152 @@ def read_batch(
  raise NotImplementedError
 
 
+class DynamicDictType(DataType, DataSerializer):
+ """
+ Dynamic DataType for dict without fixed set of keys
+ """
+
+ type: ClassVar[str] = "d_dict"
+
+ key_type: PrimitiveType
+ value_type: DataType
+
+ @validator("key_type")
+ def is_valid_key_type( # pylint: disable=no-self-argument
+ cls, key_type # noqa: B902
+ ):
+ if key_type.ptype not in ["str", "int", "float"]:
+ raise ValueError(f"key_type {key_type.ptype} is not supported")
+ return key_type
+
+ def deserialize(self, obj):
+ self.check_type(obj, dict, DeserializationError)
+ return {
+ self.key_type.get_serializer()
+ .deserialize(
+ k,
+ ): self.value_type.get_serializer()
+ .deserialize(
+ v,
+ )
+ for k, v in obj.items()
+ }
+
+ def serialize(self, instance: dict):
+ self._check_types(instance, SerializationError)
+
+ return {
+ self.key_type.get_serializer()
+ .serialize(
+ k,
+ ): self.value_type.get_serializer()
+ .serialize(
+ v,
+ )
+ for k, v in instance.items()
+ }
+
+ @classmethod
+ def create(
+ cls, obj, is_dynamic: bool = True, **kwargs
+ ) -> "DynamicDictType":
+ return DynamicDictType(
+ key_type=DataAnalyzer.analyze(
+ next(iter(obj.keys())), is_dynamic=True, **kwargs
+ ),
+ value_type=DataAnalyzer.analyze(
+ next(iter(obj.values())), is_dynamic=True, **kwargs
+ ),
+ )
+
+ def _check_types(self, obj, exc_type, ignore_key_type: bool = False):
+ self.check_type(obj, dict, exc_type)
+
+ obj_type = self.create(obj)
+ if ignore_key_type:
+ obj_types: Union[
+ Tuple[PrimitiveType, DataType], Tuple[DataType]
+ ] = (obj_type.value_type,)
+ expected_types: Union[
+ Tuple[PrimitiveType, DataType], Tuple[DataType]
+ ] = (self.value_type,)
+ else:
+ obj_types = (obj_type.key_type, obj_type.value_type)
+ expected_types = (self.key_type, self.value_type)
+ if obj_types != expected_types:
+ raise exc_type(
+ f"given dict has type: {obj_types}, expected: {expected_types}"
+ )
+
+ def get_requirements(self) -> Requirements:
+ return sum(
+ [
+ self.key_type.get_requirements(),
+ self.value_type.get_requirements(),
+ ],
+ Requirements.new(),
+ )
+
+ def get_writer(
+ self, project: str = None, filename: str = None, **kwargs
+ ) -> "DynamicDictWriter":
+ return DynamicDictWriter(**kwargs)
+
+ def get_model(self, prefix="") -> Type[BaseModel]:
+ field_type = (
+ Dict[ # type: ignore
+ self.key_type.get_serializer().get_model(
+ prefix + "_key_" # noqa: F821
+ ),
+ self.value_type.get_serializer().get_model(
+ prefix + "_val_" # noqa: F821
+ ),
+ ],
+ ...,
+ )
+ return create_model(prefix + "DynamicDictType", __root__=field_type) # type: ignore
+
+
+class DynamicDictWriter(DataWriter):
+ type: ClassVar[str] = "d_dict"
+
+ def write(
+ self, data: DataType, storage: Storage, path: str
+ ) -> Tuple[DataReader, Artifacts]:
+ if not isinstance(data, DynamicDictType):
+ raise ValueError(
+ f"expected data to be of DynamicDictTypeWriter, got {type(data)} instead"
+ )
+ with storage.open(path) as (f, art):
+ f.write(
+ json.dumps(data.get_serializer().serialize(data.data)).encode(
+ "utf-8"
+ )
+ )
+ return DynamicDictReader(data_type=data), {DataWriter.art_name: art}
+
+
+class DynamicDictReader(DataReader):
+ type: ClassVar[str] = "d_dict"
+ data_type: DynamicDictType
+
+ def read(self, artifacts: Artifacts) -> DataType:
+ if DataWriter.art_name not in artifacts:
+ raise ValueError(
+ f"Wrong artifacts {artifacts}: should be one {DataWriter.art_name} file"
+ )
+ with artifacts[DataWriter.art_name].open() as f:
+ data = json.load(f)
+ # json stores keys as strings. Deserialize string keys as well as values.
+ data = self.data_type.deserialize(data)
+ return self.data_type.copy().bind(data)
+
+ def read_batch(
+ self, artifacts: Artifacts, batch_size: int
+ ) -> Iterator[DataType]:
+ raise NotImplementedError
+
+
 #
 #
 # class BytesDataType(DataType):