Data types supported by DataChain must be of type
DataType. DataType includes most Python types
supported in Pydantic fields, as well as any class that
inherits from Pydantic BaseModel.
Pydantic models can group and nest multiple fields into one type. When reading a saved
dataset, DataChain reuses a matching model class that is already imported. If the class
is unavailable, DataChain rebuilds one from the stored schema. Models may alternatively
inherit from DataModel, a lightweight
BaseModel wrapper that registers subclasses automatically.
@staticmethoddefregister(models:DataType|Sequence[DataType]):"""For registering classes manually. It accepts a single class or a sequence of classes."""ifnotisinstance(models,Sequence):models=[models]forvalinmodels:ModelStore.register(val)
defis_chain_type(t:type)->bool:"""Return true if type is supported by `DataChain`."""ifModelStore.is_pydantic(t):returnTrueifany(tisftortisget_args(ft)[0]forftinget_args(StandardType)):returnTrueinner,is_optional=unwrap_optional(t)ifis_optional:returnis_chain_type(inner)# Deliberately not using `annotation_parts` here. This is validation, not# traversal: it must see the raw args, since normalising away `Ellipsis` would# let `list[int, ...]` through to be serialized as `list[int]`. Only `list` and# `dict` at the exact arity `type_to_str` can write back out are accepted --# abstract origins serialize as a bare "Sequence"/"Mapping", and `python_to_sql`# mis-types tuples. Matching on the origin identity also avoids `issubclass`# against generics that reject it (TypedDicts, some protocols).orig=get_origin(t)args=get_args(t)iforigislist:returnlen(args)==1andis_chain_type(args[0])iforigisdict:returnlen(args)==2andall(is_chain_type(arg)forarginargs)returnFalse