From e41fb55d182ddd0597eed37b72b3645070826f5e Mon Sep 17 00:00:00 2001 From: samsja Date: Tue, 26 Sep 2023 12:50:19 +0200 Subject: [PATCH] fix: make DocList compatible with BaseDocWithoutId Signed-off-by: samsja --- docarray/array/any_array.py | 14 +++++++------- docarray/array/doc_list/doc_list.py | 17 +++++++++-------- tests/units/document/test_doc_wo_id.py | 11 +++++++++++ 3 files changed, 27 insertions(+), 15 deletions(-) create mode 100644 tests/units/document/test_doc_wo_id.py diff --git a/docarray/array/any_array.py b/docarray/array/any_array.py index 1b92f01f72..28373e2306 100644 --- a/docarray/array/any_array.py +++ b/docarray/array/any_array.py @@ -19,7 +19,7 @@ import numpy as np -from docarray.base_doc import BaseDoc +from docarray.base_doc.doc import BaseDocWithoutId from docarray.display.document_array_summary import DocArraySummary from docarray.exceptions.exceptions import UnusableObjectError from docarray.typing.abstract_type import AbstractType @@ -30,7 +30,7 @@ from docarray.typing.tensor.abstract_tensor import AbstractTensor T = TypeVar('T', bound='AnyDocArray') -T_doc = TypeVar('T_doc', bound=BaseDoc) +T_doc = TypeVar('T_doc', bound=BaseDocWithoutId) IndexIterType = Union[slice, Iterable[int], Iterable[bool], None] UNUSABLE_ERROR_MSG = ( @@ -42,18 +42,18 @@ class AnyDocArray(Sequence[T_doc], Generic[T_doc], AbstractType): - doc_type: Type[BaseDoc] - __typed_da__: Dict[Type['AnyDocArray'], Dict[Type[BaseDoc], Type]] = {} + doc_type: Type[BaseDocWithoutId] + __typed_da__: Dict[Type['AnyDocArray'], Dict[Type[BaseDocWithoutId], Type]] = {} def __repr__(self): return f'<{self.__class__.__name__} (length={len(self)})>' @classmethod - def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]): + def __class_getitem__(cls, item: Union[Type[BaseDocWithoutId], TypeVar, str]): if not isinstance(item, type): return Generic.__class_getitem__.__func__(cls, item) # type: ignore # this do nothing that checking that item is valid type var or str - if not safe_issubclass(item, BaseDoc): + if not safe_issubclass(item, BaseDocWithoutId): raise ValueError( f'{cls.__name__}[item] item should be a Document not a {item} ' ) @@ -66,7 +66,7 @@ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]): global _DocArrayTyped class _DocArrayTyped(cls): # type: ignore - doc_type: Type[BaseDoc] = cast(Type[BaseDoc], item) + doc_type: Type[BaseDocWithoutId] = cast(Type[BaseDocWithoutId], item) for field in _DocArrayTyped.doc_type._docarray_fields().keys(): diff --git a/docarray/array/doc_list/doc_list.py b/docarray/array/doc_list/doc_list.py index f27e31d174..2057a0ace9 100644 --- a/docarray/array/doc_list/doc_list.py +++ b/docarray/array/doc_list/doc_list.py @@ -16,13 +16,14 @@ from pydantic import parse_obj_as from typing_extensions import SupportsIndex -from typing_inspect import is_union_type, is_typevar +from typing_inspect import is_typevar, is_union_type from docarray.array.any_array import AnyDocArray from docarray.array.doc_list.io import IOMixinDocList from docarray.array.doc_list.pushpull import PushPullMixin from docarray.array.list_advance_indexing import IndexIterType, ListAdvancedIndexing -from docarray.base_doc import AnyDoc, BaseDoc +from docarray.base_doc import AnyDoc +from docarray.base_doc.doc import BaseDocWithoutId from docarray.typing import NdArray from docarray.utils._internal.pydantic import is_pydantic_v2 @@ -40,7 +41,7 @@ from docarray.typing.tensor.abstract_tensor import AbstractTensor T = TypeVar('T', bound='DocList') -T_doc = TypeVar('T_doc', bound=BaseDoc) +T_doc = TypeVar('T_doc', bound=BaseDocWithoutId) class DocList( @@ -120,7 +121,7 @@ class Image(BaseDoc): """ - doc_type: Type[BaseDoc] = AnyDoc + doc_type: Type[BaseDocWithoutId] = AnyDoc def __init__( self, @@ -229,7 +230,7 @@ def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]): not is_union_type(field_type) and is_field_required and isinstance(field_type, type) - and safe_issubclass(field_type, BaseDoc) + and safe_issubclass(field_type, BaseDocWithoutId) ): # calling __class_getitem__ ourselves is a hack otherwise mypy complain # most likely a bug in mypy though @@ -273,7 +274,7 @@ def to_doc_vec( @classmethod def _docarray_validate( cls: Type[T], - value: Union[T, Iterable[BaseDoc]], + value: Union[T, Iterable[BaseDocWithoutId]], ): from docarray.array.doc_vec.doc_vec import DocVec @@ -333,9 +334,9 @@ def __getitem__(self, item): return super().__getitem__(item) @classmethod - def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]): + def __class_getitem__(cls, item: Union[Type[BaseDocWithoutId], TypeVar, str]): - if isinstance(item, type) and safe_issubclass(item, BaseDoc): + if isinstance(item, type) and safe_issubclass(item, BaseDocWithoutId): return AnyDocArray.__class_getitem__.__func__(cls, item) # type: ignore if ( isinstance(item, object) diff --git a/tests/units/document/test_doc_wo_id.py b/tests/units/document/test_doc_wo_id.py new file mode 100644 index 0000000000..78172a0387 --- /dev/null +++ b/tests/units/document/test_doc_wo_id.py @@ -0,0 +1,11 @@ +from docarray import DocList +from docarray.base_doc.doc import BaseDocWithoutId + + +def test_doc_list(): + class A(BaseDocWithoutId): + text: str + + cls_doc_list = DocList[A] + + assert isinstance(cls_doc_list, type)