docarray · JohannesMessner · Nov 17, 2022 · Nov 16, 2022 · Nov 16, 2022 · Nov 16, 2022
diff --git a/docarray/array/abstract_array.py b/docarray/array/abstract_array.py
@@ -2,13 +2,18 @@
 from typing import Iterable, Type
 
 from docarray.document import BaseDocument
-from docarray.document.abstract_document import AbstractDocument
 
 
 class AbstractDocumentArray(Iterable):
 
     document_type: Type[BaseDocument]
 
     @abstractmethod
-    def __init__(self, docs: Iterable[AbstractDocument]):
+    def __init__(self, docs: Iterable[BaseDocument]):
+        ...
+
+    @abstractmethod
+    def __class_getitem__(
+        cls, item: Type[BaseDocument]
+    ) -> Type['AbstractDocumentArray']:
         ...
diff --git a/docarray/array/array.py b/docarray/array/array.py
@@ -3,7 +3,6 @@
 from docarray.array.abstract_array import AbstractDocumentArray
 from docarray.array.mixins import GetAttributeArrayMixin, ProtoArrayMixin
 from docarray.document import AnyDocument, BaseDocument, BaseNode
-from docarray.document.abstract_document import AbstractDocument
 
 
 class DocumentArray(
@@ -21,7 +20,7 @@ class DocumentArray(
 
     document_type: Type[BaseDocument] = AnyDocument
 
-    def __init__(self, docs: Iterable[AbstractDocument]):
+    def __init__(self, docs: Iterable[BaseDocument]):
         super().__init__(doc_ for doc_ in docs)
 
     def __class_getitem__(cls, item: Type[BaseDocument]):
@@ -31,7 +30,7 @@ def __class_getitem__(cls, item: Type[BaseDocument]):
             )
 
         class _DocumenArrayTyped(DocumentArray):
-            document_type = item
+            document_type: Type[BaseDocument] = item
 
         for field in _DocumenArrayTyped.document_type.__fields__.keys():
 

diff --git a/docarray/array/mixins/attribute.py b/docarray/array/mixins/attribute.py
@@ -1,17 +1,30 @@
-from typing import List
+from typing import List, Union
 
 from docarray.array.abstract_array import AbstractDocumentArray
+from docarray.document import BaseDocument
 
 
 class GetAttributeArrayMixin(AbstractDocumentArray):
     """Helpers that provide attributes getter in bulk"""
 
-    def _get_documents_attribute(self, field: str) -> List:
+    def _get_documents_attribute(
+        self, field: str
+    ) -> Union[List, AbstractDocumentArray]:
         """Return all values of the fields from all docs this array contains
 
         :param field: name of the fields to extract
         :return: Returns a list of the field value for each document
         in the array like container
         """
 
-        return [getattr(doc, field) for doc in self]
+        field_type = self.__class__.document_type._get_nested_document_class(field)
+
+        if issubclass(field_type, BaseDocument):
+            # calling __class_getitem__ ourselves is a hack otherwise mypy complain
+            # most likely a bug in mypy though
+            # bug reported here https://github.com/python/mypy/issues/14111
+            return self.__class__.__class_getitem__(field_type)(
+                (getattr(doc, field) for doc in self)
+            )
+        else:
+            return [getattr(doc, field) for doc in self]
diff --git a/docarray/document/mixins/proto.py b/docarray/document/mixins/proto.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Type, TypeVar
 
 from pydantic.tools import parse_obj_as
 
@@ -7,10 +7,12 @@
 from docarray.proto import DocumentProto, NodeProto
 from docarray.typing import ID, AnyUrl, Embedding, ImageUrl, Tensor, TorchTensor
 
+T = TypeVar('T', bound='ProtoMixin')
+
 
 class ProtoMixin(AbstractDocument, BaseNode):
     @classmethod
-    def from_protobuf(cls, pb_msg: 'DocumentProto') -> 'ProtoMixin':
+    def from_protobuf(cls: Type[T], pb_msg: 'DocumentProto') -> T:
         """create a Document from a protobuf message"""
         from docarray import DocumentArray
 

diff --git a/tests/units/array/test_mixins/test_attribute.py b/tests/units/array/test_mixins/test_attribute.py
@@ -51,3 +51,19 @@ class Mmdoc(BaseDocument):
     assert len(texts) == N
     for i, text in enumerate(texts):
         assert text == f'hello{i}'
+
+
+def test_get_bulk_attributes_document():
+    class InnerDoc(BaseDocument):
+        text: str
+
+    class Mmdoc(BaseDocument):
+        inner: InnerDoc
+
+    N = 10
+
+    da = DocumentArray[Mmdoc](
+        (Mmdoc(inner=InnerDoc(text=f'hello{i}')) for i in range(N))
+    )
+
+    assert isinstance(da.inner, DocumentArray)