From fc41eb01eaaab0cc16dd99451d1e01f51d16a9f1 Mon Sep 17 00:00:00 2001 From: Joan Fontanals Martinez Date: Tue, 19 Sep 2023 23:53:15 +0200 Subject: [PATCH] fix: fix double subscriptable error Signed-off-by: Joan Fontanals Martinez --- docarray/array/doc_list/doc_list.py | 2 ++ tests/units/array/test_array.py | 16 +++++++++++++--- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/docarray/array/doc_list/doc_list.py b/docarray/array/doc_list/doc_list.py index f27e31d174..82e00ce121 100644 --- a/docarray/array/doc_list/doc_list.py +++ b/docarray/array/doc_list/doc_list.py @@ -334,6 +334,8 @@ def __getitem__(self, item): @classmethod def __class_getitem__(cls, item: Union[Type[BaseDoc], TypeVar, str]): + if cls.doc_type != AnyDoc: + raise TypeError(f'{cls} object is not subscriptable') if isinstance(item, type) and safe_issubclass(item, BaseDoc): return AnyDocArray.__class_getitem__.__func__(cls, item) # type: ignore diff --git a/tests/units/array/test_array.py b/tests/units/array/test_array.py index ab078ebac7..6c2308009c 100644 --- a/tests/units/array/test_array.py +++ b/tests/units/array/test_array.py @@ -467,7 +467,6 @@ class Image(BaseDoc): def test_validate_list_dict(): - images = [ dict(url=f'http://url.com/foo_{i}.png', tensor=NdArray(i)) for i in [2, 0, 1] ] @@ -493,7 +492,18 @@ def test_parameterize_list(): from docarray import DocList, BaseDoc with pytest.raises(TypeError) as excinfo: - doc = DocList[BaseDoc()] - assert doc is None + da = DocList[BaseDoc()] + assert da is None assert str(excinfo.value) == 'Expecting a type, got object instead' + + +def test_not_double_subcriptable(): + from docarray import DocList + from docarray.documents import TextDoc + + with pytest.raises(TypeError) as excinfo: + da = DocList[TextDoc][TextDoc] + assert da is None + + assert 'not subscriptable' in str(excinfo.value)