text: import heapq
  import itertools
  from collections.abc import MutableSequence
  from typing import (
      Union,
      Iterable,
      Iterator,
      Optional,
      Generator,
      Dict,
      Callable,
      List,
  )
  
  from .mixins import AllMixins
  from .types import DocumentArraySourceType
  from ..document import Document
  from ...helper import typename
  from ...proto import jina_pb2
  
  try:
      # when protobuf using Cpp backend
      from google.protobuf.pyext._message import (
          RepeatedCompositeContainer as RepeatedContainer,
      )
  except:
      # when protobuf using Python backend
      from google.protobuf.internal.containers import (
          RepeatedCompositeFieldContainer as RepeatedContainer,
      )
  
  __all__ = ['DocumentArray']
  
  
  class DocumentArray(
      AllMixins,
      MutableSequence,
  ):
      """
      :class:`DocumentArray` is a mutable sequence of :class:`Document`.
      It gives an efficient view of a list of Document. One can iterate over it like
      a generator but ALSO modify it, count it, get item, or union two 'DocumentArray's using the '+' and '+=' operators.
  
      It is supposed to act as a view containing a pointer to a `RepeatedContainer` of `DocumentProto` while offering `Document` Jina native types
      when getting items or iterating over it
  
      :param docs: the document array to construct from. One can also give `DocumentArrayProto` directly, then depending on the ``copy``,
                  it builds a view or a copy from it. It also can accept a List
      """
  
      def __init__(self, docs: Optional[DocumentArraySourceType] = None):
          super().__init__()
          self._pb_body = []
          if docs is not None:
              if isinstance(docs, jina_pb2.DocumentArrayProto):
                  # This would happen when loading from file or memmap
                  self._pb_body = docs.docs
              elif isinstance(docs, RepeatedContainer):
                  # This would happen when `doc.matches` or `doc.chunks`
                  self._pb_body = docs
              elif isinstance(docs, DocumentArray):
                  # This would happen in the client
                  self._pb_body = docs._pb_body
              else:
                  if isinstance(docs, Document):
                      # single Document
                      docs = [docs]
  
                  from .memmap import DocumentArrayMemmap
  
                  if isinstance(
                      docs, (list, tuple, Generator, DocumentArrayMemmap, itertools.chain)
                  ):
                      # This would happen in the client
                      for doc in docs:
                          if isinstance(doc, Document):
                              self._pb_body.append(doc.proto)
                          elif isinstance(doc, jina_pb2.DocumentProto):
                              self._pb_body.append(doc)
                          else:
                              raise ValueError(f'Unexpected element in an input list')
                  else:
                      raise ValueError(
                          f'DocumentArray got an unexpected input {type(docs)}'
                      )
          self._id_to_index = None
  
      @property
      def _index_map(self) -> Dict:
          """Return the `_id_to_index` map
  
          :return: a Python dict.
          """
          if not self._id_to_index:
              self._update_id_to_index_map()
          return self._id_to_index
  
      def _update_id_to_index_map(self):
          """Update the id_to_index map by enumerating all Documents in self._pb_body.
  
          Very costy! Only use this function when self._pb_body is dramtically changed.
          """
  
          self._id_to_index = {
              d.id: i for i, d in enumerate(self._pb_body)
          }  # type: Dict[str, int]
  
      def insert(self, index: int, doc: 'Document') -> None:
          """Insert `doc` at `index`.
  
          :param index: Position of the insertion.
          :param doc: The doc needs to be inserted.
          """
          self._pb_body.insert(index, doc.proto)
          if self._id_to_index:
              self._id_to_index[doc.id] = index
  
      def __setitem__(self, key, value: 'Document'):
          if isinstance(key, int):
              self[key].CopyFrom(value)
              if self._id_to_index:
                  self._id_to_index[value.id] = key
          elif isinstance(key, str):
              self[self._index_map[key]].CopyFrom(value)
          else:
              raise IndexError(f'do not support this index {key}')
  
      def __delitem__(self, index: Union[int, str, slice]):
          if isinstance(index, int):
              del self._pb_body[index]
          elif isinstance(index, str):
              del self[self._index_map[index]]
              self._index_map.pop(index)
          elif isinstance(index, slice):
              del self._pb_body[index]
          else:
              raise IndexError(
                  f'do not support this index type {typename(index)}: {index}'
              )
  
      def __eq__(self, other):
          return (
              type(self) is type(other)
              and type(self._pb_body) is type(other._pb_body)
              and self._pb_body == other._pb_body
          )
  
      def __len__(self):
          return len(self._pb_body)
  
      def __iter__(self) -> Iterator['Document']:
          for d in self._pb_body:
              yield Document(d)
  
      def __contains__(self, item: str):
          return item in self._index_map
  
      def __getitem__(self, item: Union[int, str, slice, List]):
          if isinstance(item, int):
              return Document(self._pb_body[item])
          elif isinstance(item, str):
              return self[self._index_map[item]]
          elif isinstance(item, slice):
              return DocumentArray(self._pb_body[item])
          elif isinstance(item, list):
              return DocumentArray(self._pb_body[t] for t in item)
          else:
              raise IndexError(f'do not support this index type {typename(item)}: {item}')
  
      def append(self, doc: 'Document'):
          """
          Append `doc` in :class:`DocumentArray`.
  
          :param doc: The doc needs to be appended.
          """
          if self._id_to_index:
              self._id_to_index[doc.id] = len(self._pb_body)
          self._pb_body.append(doc.proto)
  
      def extend(self, docs: Iterable['Document']) -> None:
          """
          Extend the :class:`DocumentArray` by appending all the items from the iterable.
  
          :param docs: the iterable of Documents to extend this array with
          """
          if not docs:
              return
  
          for doc in docs:
              self.append(doc)
  
      def clear(self):
          """Clear the data of :class:`DocumentArray`"""
          del self._pb_body[:]
          if self._id_to_index:
              self._id_to_index.clear()
  
      def reverse(self):
          """In-place reverse the sequence."""
          size = len(self._pb_body)
          hi_idx = size - 1
          for i in range(int(size / 2)):
              tmp = jina_pb2.DocumentProto()
              tmp.CopyFrom(self._pb_body[hi_idx])
              self._pb_body[hi_idx].CopyFrom(self._pb_body[i])
              self._pb_body[i].CopyFrom(tmp)
              hi_idx -= 1
          self._update_id_to_index_map()
  
      def sort(
          self,
          key: Callable,
          top_k: Optional[int] = None,
          reverse: bool = False,
      ):
          """
          Sort the items of the :class:`DocumentArray` in place.
  
          :param key: key callable to sort based upon
          :param top_k: make sure that the first `topk` elements are correctly sorted rather than
              sorting the entire list
          :param reverse: reverse=True will sort the list in descending order. Default is False
          """
  
          def overriden_key(proto):
              # Function to override the `proto` and wrap it around a `Document` to enable sorting via
              # `Document-like` interface
              d = Document(proto)
              return key(d)
  
          # Logic here: `overriden_key` is offered to allow the user sort via pythonic `Document` syntax. However,
          # maybe there may be cases where this won't work and the user may enter `proto-like` interface. To make
          # sure (quite fragile) the `sort` will work seamlessly, it tries to apply `key` to the first element and
          # see if it works. If it works it can sort with `proto` interface, otherwise use `Document` interface one.
          # (Very often the 2 interfaces are both the same and valid, so proto will have less overhead
          _key = key
          try:
              key(self._pb_body[0])
          except:
              _key = overriden_key
  
          if top_k is None or top_k >= len(self._pb_body):
              self._pb_body.sort(key=_key, reverse=reverse)
          else:
              # heap based sorting
              # makes sure that the top_k elements are correctly sorted and leaves the rest unsorted
              if _key is None:
                  # heap is a list of documents
                  heap = [element for element in self._pb_body]
              else:
                  # heap is a list of tuples (key, document)
                  heap = [
                      (_key(element), i, element)
                      for i, element in enumerate(self._pb_body)
                  ]
              # if reverse use the maxheap operations for .heapify and .heappop
              heapify = heapq._heapify_max if reverse else heapq.heapify
              heappop = heapq._heappop_max if reverse else heapq.heappop
  
              # transform the original list to a heap and pop the top k elements
              heapify(heap)
              topk = [heappop(heap) for _ in range(top_k)]
  
              # get back to lists of docs from the lists of tuples
              _, _, topk = zip(*topk)
              _, _, heap = zip(*heap)
              topk, heap = list(topk), list(heap)
              # update the protobuf body
              self._pb_body = topk + heap
  
          self._update_id_to_index_map()
  
      @staticmethod
      def _flatten(sequence) -> 'DocumentArray':
          return DocumentArray(list(itertools.chain.from_iterable(sequence)))
  