style: reformat code (#227)

docarray/__init__.py

@@ -5,6 +5,7 @@
 from .document import Document
 from .array import DocumentArray
 
+
 if 'DA_NO_RICH_HANDLER' not in os.environ:
     from rich.traceback import install
 

docarray/array/base.py

@@ -1,4 +1,4 @@
-from typing import MutableSequence, TYPE_CHECKING, Union, Iterable, Type
+from typing import MutableSequence, TYPE_CHECKING, Union, Iterable
 
 from .. import Document
 

docarray/array/document.py

@@ -58,7 +58,7 @@ def __new__(
     def __new__(
         cls,
         _docs: Optional['DocumentArraySourceType'] = None,
-        storage: str = 'elastic',
+        storage: str = 'elasticsearch',
         config: Optional[Union['ElasticConfig', Dict]] = None,
     ) -> 'DocumentArrayElastic':
         """Create a Elastic-powered DocumentArray object."""
@@ -86,7 +86,7 @@ def __new__(cls, *args, storage: str = 'memory', **kwargs):
                 from .qdrant import DocumentArrayQdrant
 
                 instance = super().__new__(DocumentArrayQdrant)
-            elif storage == 'elastic':
+            elif storage == 'elasticsearch':
                 from .elastic import DocumentArrayElastic
 
                 instance = super().__new__(DocumentArrayElastic)

docarray/array/elastic.py

@@ -5,7 +5,7 @@
 
 
 class DocumentArrayElastic(StorageMixins, DocumentArray):
-    """This is a :class:`DocumentArray` that uses Elastic Search as
+    """This is a :class:`DocumentArray` that uses Elasticsearch as
     vector search engine and storage.
     """
 

docarray/array/mixins/io/binary.py

@@ -1,13 +1,12 @@
 import base64
 import io
-import os.path
 import os
+import os.path
 import pickle
 from contextlib import nullcontext
 from pathlib import Path
 from typing import Union, BinaryIO, TYPE_CHECKING, Type, Optional, Generator
 
-
 from ....helper import (
     get_compress_ctx,
     decompress_bytes,

docarray/array/mixins/io/json.py

@@ -1,5 +1,4 @@
 import json
-import os.path
 from contextlib import nullcontext
 from typing import Union, TextIO, TYPE_CHECKING, Type, List
 

docarray/array/mixins/plot.py

@@ -29,6 +29,8 @@ def summary(self):
 
         from rich.table import Table
         from rich.console import Console
+        from rich.panel import Panel
+
         from rich import box
 
         tables = []
@@ -37,7 +39,7 @@ def summary(self):
         all_attrs = self._get_attributes('non_empty_fields')
         attr_counter = Counter(all_attrs)
 
-        table = Table(box=box.SIMPLE, title='Documents Summary')
+        table = Table(box=box.SIMPLE)
         table.show_header = False
         table.add_row('Length', str(len(self)))
         is_homo = len(attr_counter) == 1
@@ -70,12 +72,12 @@ def summary(self):
                     _text = f'{_doc_text} attributes'
                 table.add_row(_text, str(_a))
 
-        tables.append(table)
+        tables.append(Panel(table, title='Documents Summary', expand=False))
 
         all_attrs_names = tuple(sorted(all_attrs_names))
         if all_attrs_names:
 
-            attr_table = Table(box=box.SIMPLE, title='Attributes Summary')
+            attr_table = Table(box=box.SIMPLE)
             attr_table.add_column('Attribute')
             attr_table.add_column('Data type')
             attr_table.add_column('#Unique values')
@@ -96,16 +98,16 @@ def summary(self):
                     str(len(_a)),
                     str(any(_aa is None for _aa in _a)),
                 )
-            tables.append(attr_table)
+            tables.append(Panel(attr_table, title='Attributes Summary', expand=False))
 
         storage_infos = self._get_storage_infos()
         if storage_infos:
-            storage_table = Table(box=box.SIMPLE, title='Storage Summary')
+            storage_table = Table(box=box.SIMPLE)
             storage_table.show_header = False
             for k, v in storage_infos.items():
                 storage_table.add_row(k, v)
 
-            tables.append(storage_table)
+            tables.append(Panel(storage_table, title='Storage Summary', expand=False))
 
         console.print(*tables)
 

docarray/array/mixins/traverse.py

@@ -5,7 +5,6 @@
     TYPE_CHECKING,
     Optional,
     Callable,
-    Tuple,
     Dict,
     List,
 )
@@ -14,7 +13,6 @@
     from ... import DocumentArray, Document
     from ...types import T
 
-
 ATTRIBUTES_SEPARATOR = ','
 PATHS_SEPARATOR = ','
 
@@ -41,7 +39,6 @@
 TRAVERSAL_PATH = rf'{SELECTOR}{SLICE}{REMAINDER}'
 TRAVERSAL_PATH_TAGGED = rf'(?P<path>{SELECTOR_TAGGED}{SLICE_TAGGED}){REMAINDER_TAGGED}'
 
-
 PATHS_REMAINDER_TAGGED = rf'(?P<paths_remainder>({PATHS_SEPARATOR}{TRAVERSAL_PATH})*)'
 
 TRAVERSAL_PATH_LIST_TAGGED = (

docarray/array/qdrant.py

@@ -1,12 +1,8 @@
-from typing import Iterable
-
 from .document import DocumentArray
 from .storage.qdrant import StorageMixins, QdrantConfig
 
 __all__ = ['DocumentArrayQdrant', 'QdrantConfig']
 
-from .. import Document
-
 
 class DocumentArrayQdrant(StorageMixins, DocumentArray):
     """This is a :class:`DocumentArray` that uses Qdrant as

docarray/array/queryset/lookup.py

@@ -26,7 +26,6 @@
 WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 """
-import re
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:

docarray/array/storage/annlite/backend.py

@@ -1,4 +1,3 @@
-import numpy as np
 from dataclasses import dataclass, asdict, field
 from typing import (
     Union,
@@ -8,6 +7,8 @@
     Iterable,
 )
 
+import numpy as np
+
 from ..base.backend import BaseBackendMixin
 from ....helper import dataclass_from_dict
 
@@ -64,7 +65,6 @@ def _init_storage(
         from annlite import AnnLite
 
         self._annlite = AnnLite(self.n_dim, lock=False, **config)
-        from ... import DocumentArray
         from .... import Document
 
         super()._init_storage()

docarray/array/storage/annlite/seqlike.py

@@ -1,9 +1,8 @@
-from typing import Union, Iterable, Sequence
+from typing import Union, Iterable
 
 from ..base.seqlike import BaseSequenceLikeMixin
-from .... import Document
-
 from ...memory import DocumentArrayInMemory
+from .... import Document
 
 
 class SequenceLikeMixin(BaseSequenceLikeMixin):

docarray/array/storage/base/seqlike.py

@@ -1,5 +1,5 @@
-from typing import Iterator, Iterable, MutableSequence
 from abc import abstractmethod
+from typing import Iterator, Iterable, MutableSequence
 
 from .... import Document
 

docarray/array/storage/elastic/__init__.py

@@ -1,8 +1,8 @@
 from abc import ABC
 
 from .backend import BackendMixin, ElasticConfig
-from .getsetdel import GetSetDelMixin
 from .find import FindMixin
+from .getsetdel import GetSetDelMixin
 from .seqlike import SequenceLikeMixin
 
 __all__ = ['StorageMixins', 'ElasticConfig']

docarray/array/storage/elastic/backend.py

@@ -1,56 +1,38 @@
 import copy
-import itertools
-import warnings
+import uuid
 from dataclasses import dataclass, field
 from typing import (
     Dict,
     Optional,
     TYPE_CHECKING,
     Union,
-    Tuple,
     List,
-    Sequence,
-    Generator,
-    Iterator,
     Iterable,
+    Any,
 )
 
 import numpy as np
-import uuid
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import bulk
 
 from ..base.backend import BaseBackendMixin
-from .... import DocumentArray, Document
+from .... import Document
 from ....helper import dataclass_from_dict
 
-from elasticsearch import Elasticsearch
-from elasticsearch.helpers import bulk
-
 if TYPE_CHECKING:
     from ....types import (
         DocumentArraySourceType,
     )
     from ....types import DocumentArraySourceType, ArrayType
 
-from docarray.math.helper import EPSILON
-
-
-def _sanitize_table_name(table_name: str) -> str:
-    ret = ''.join(c for c in table_name if c.isalnum() or c == '_')
-    if ret != table_name:
-        warnings.warn(f'The table name is changed to {ret} due to illegal characters')
-    return ret
-
 
 @dataclass
 class ElasticConfig:
     n_dim: int  # dims  in elastic
-    basic_auth: Optional[Tuple[str, str]] = None
-    ca_certs: Optional[str] = None
     distance: str = 'cosine'  # similarity in elastic
-    host: Optional[str] = field(default='http://localhost')
-    port: Optional[int] = field(default=9200)
-    index_name: Optional[str] = field(default=None)
-    serialize_config: Dict = field(default_factory=dict)
+    hosts: str = 'http://localhost:9200'
+    index_name: Optional[str] = None
+    es_config: Dict[str, Any] = field(default_factory=dict)
 
 
 class BackendMixin(BaseBackendMixin):
@@ -97,9 +79,6 @@ def _build_offset2id_index(self):
         if not self._client.indices.exists(index=self._index_name_offset2id):
             self._client.indices.create(index=self._index_name_offset2id, ignore=[404])
 
-    def _build_hosts(self):
-        return self._config.host + ':' + str(self._config.port)
-
     def _build_schema_from_elastic_config(self, elastic_config):
         da_schema = {
             "mappings": {
@@ -120,9 +99,8 @@ def _build_schema_from_elastic_config(self, elastic_config):
     def _build_client(self):
 
         client = Elasticsearch(
-            hosts=self._build_hosts(),
-            ca_certs=self._config.ca_certs,
-            basic_auth=self._config.basic_auth,
+            hosts=self._config.hosts,
+            **self._config.es_config,
         )
 
         schema = self._build_schema_from_elastic_config(self._config)
@@ -147,9 +125,9 @@ def _doc_id_exists(self, doc_id):
     def _get_storage_infos(self) -> Dict:
         return {
             'Backend': 'ElasticSearch',
-            'Host': str(self._config.host),
-            'Port': str(self._config.port),
-            'distance': str(self._config.distance),
+            'Hosts': str(self._config.hosts),
+            'ES config': str(self._config.es_config),
+            'Distance': str(self._config.distance),
             'Vector dimension': str(self._config.n_dim),
         }
 
@@ -189,6 +167,8 @@ def _get_offset2ids_meta(self) -> List:
             return []
 
     def _map_embedding(self, embedding: 'ArrayType') -> List[float]:
+        from ....math.helper import EPSILON
+
         if embedding is None:
             embedding = np.zeros(self.n_dim) + EPSILON
         else:

docarray/array/storage/elastic/getsetdel.py

@@ -1,4 +1,3 @@
-import copy
 from typing import Iterable, Dict
 
 from ..base.getsetdel import BaseGetSetDelMixin

docarray/array/storage/elastic/seqlike.py

@@ -2,7 +2,6 @@
 
 from ..base.seqlike import BaseSequenceLikeMixin
 from .... import Document
-from ..registry import _REGISTRY
 
 
 class SequenceLikeMixin(BaseSequenceLikeMixin):

docarray/array/storage/memory/backend.py

@@ -25,8 +25,6 @@ def _init_storage(
     ):
         super()._init_storage(_docs, copy=copy, *args, **kwargs)
 
-        from ... import DocumentArray
-
         self._data = {}
         if _docs is None:
             return

docarray/array/storage/memory/seqlike.py

@@ -1,4 +1,4 @@
-from typing import Iterator, Union, Iterable
+from typing import Union, Iterable
 
 from ..base.seqlike import BaseSequenceLikeMixin
 from .... import Document

docarray/array/storage/qdrant/__init__.py

@@ -36,7 +36,6 @@ def collection_name(self) -> str:
 
     @property
     def collection_name_meta(self) -> str:
-
         return f'{self.collection_name}_meta'
 
     @property

docarray/array/storage/qdrant/backend.py

@@ -59,8 +59,6 @@ def _init_storage(
             raise an error if both are provided
         """
 
-        from ... import DocumentArray
-
         self._schemas = None
 
         if not config:

docarray/array/storage/qdrant/find.py

@@ -1,16 +1,16 @@
 from abc import abstractmethod
 from typing import (
-    Union,
     TYPE_CHECKING,
     TypeVar,
     Sequence,
     List,
 )
 
+from qdrant_openapi_client.models.models import Distance
+
 from .... import Document, DocumentArray
 from ....math import ndarray
 from ....score import NamedScore
-from qdrant_openapi_client.models.models import Distance
 
 if TYPE_CHECKING:
     import tensorflow

docarray/array/storage/qdrant/seqlike.py

@@ -1,10 +1,10 @@
 from abc import abstractmethod
 from typing import Iterable, Union
-from docarray import Document
 
 from qdrant_client import QdrantClient
 
-from docarray.array.storage.base.seqlike import BaseSequenceLikeMixin
+from ..base.seqlike import BaseSequenceLikeMixin
+from .... import Document
 
 
 class SequenceLikeMixin(BaseSequenceLikeMixin):