additional typing support

Author: miguelgrinberg

elasticsearch_dsl/__init__.py

@@ -19,7 +19,7 @@
 from .aggs import A
 from .analysis import analyzer, char_filter, normalizer, token_filter, tokenizer
 from .document import AsyncDocument, Document
-from .document_base import InnerDoc, MetaField
+from .document_base import InnerDoc, M, MetaField, mapped_field
 from .exceptions import (
     ElasticsearchDslException,
     IllegalOperation,
@@ -148,6 +148,7 @@
     "Keyword",
     "Long",
     "LongRange",
+    "M",
     "Mapping",
     "MetaField",
     "MultiSearch",
@@ -178,6 +179,7 @@
     "char_filter",
     "connections",
     "construct_field",
+    "mapped_field",
     "normalizer",
     "token_filter",
     "tokenizer",

elasticsearch_dsl/document_base.py

@@ -17,7 +17,7 @@
 
 from datetime import date, datetime
 from fnmatch import fnmatch
-from typing import List, Optional
+from typing import TYPE_CHECKING, Any, Generic, List, Optional, TypeVar, Union, overload
 
 from .exceptions import ValidationException
 from .field import (
@@ -70,6 +70,34 @@ def __init__(self, name, bases, attrs):
         # create the mapping instance
         self.mapping = getattr(meta, "mapping", Mapping())
 
+        # register the document's fields, which can be given in a few formats:
+        #
+        # class MyDocument(Document):
+        #     # required field using native typing
+        #     # (str, int, float, bool, datetime, date)
+        #     field1: str
+        #
+        #     # optional field using native typing
+        #     field2: Optional[datetime]
+        #
+        #     # array field using native typing
+        #     field3: list[int]
+        #
+        #     # sub-object, same as Object(MyInnerDoc)
+        #     field4: MyInnerDoc
+        #
+        #     # nested sub-objects, same as Nested(MyInnerDoc)
+        #     field5: list[MyInnerDoc]
+        #
+        #     # use typing, but override with any stock or custom field
+        #     field6: bool = MyCustomField()
+        #
+        #     # best mypy and pyright typing support
+        #     field7: M[date]
+        #     field8: M[str] = mapped_field(MyCustomText())
+        #
+        #     # legacy format without Python typing
+        #     field8 = Text()
         annotations = attrs.get("__annotations__", {})
         fields = set([n for n in attrs if isinstance(attrs[n], Field)])
         fields.update(annotations.keys())
@@ -81,25 +109,34 @@ def __init__(self, name, bases, attrs):
                 required = True
                 multi = False
                 while hasattr(type_, "__origin__"):
-                    if type_.__origin__ == Optional:
-                        required = False
+                    if type_.__origin__ == Mapped:
                         type_ = type_.__args__[0]
-                    elif issubclass(type_.__origin__, List):
+                    elif type_.__origin__ == Union:
+                        if len(type_.__args__) == 2 and type_.__args__[1] is type(None):
+                            required = False
+                            type_ = type_.__args__[0]
+                        else:
+                            raise TypeError("Unsupported union")
+                    elif type_.__origin__ in [list, List]:
                         multi = True
                         type_ = type_.__args__[0]
-                if issubclass(type_, InnerDoc):
+                    else:
+                        break
+                field_args = []
+                field_kwargs = {}
+                if not isinstance(type_, type):
+                    raise TypeError(f"Cannot map type {type_}")
+                elif issubclass(type_, InnerDoc):
                     field = Nested if multi else Object
-                    field_args = {}
+                    field_args = [type_]
                 elif type_ in self.type_annotation_map:
-                    field, field_args = self.type_annotation_map[type_]
+                    field, field_kwargs = self.type_annotation_map[type_]
                 elif not issubclass(type_, Field):
                     raise TypeError(f"Cannot map type {type_}")
                 else:
                     field = type_
-                    field_args = {}
-                field_args = {"multi": multi, "required": required, **field_args}
-                value = field(**field_args)
-            value._name = name
+                field_kwargs = {"multi": multi, "required": required, **field_kwargs}
+                value = field(*field_args, **field_kwargs)
             self.mapping.field(name, value)
             if name in attrs:
                 del attrs[name]
@@ -120,6 +157,36 @@ def name(self):
         return self.mapping.properties.name
 
 
+_FieldType = TypeVar("_FieldType")
+
+
+class Mapped(Generic[_FieldType]):
+    __slots__ = {}
+
+    if TYPE_CHECKING:
+
+        @overload
+        def __get__(self, instance: None, owner: Any) -> InstrumentedField: ...
+
+        @overload
+        def __get__(self, instance: object, owner: Any) -> _FieldType: ...
+
+        def __get__(
+            self, instance: Optional[object], owner: Any
+        ) -> Union[InstrumentedField, _FieldType]: ...
+
+        def __set__(self, instance: Optional[object], value: _FieldType) -> None: ...
+
+        def __delete__(self, instance: Any) -> None: ...
+
+
+M = Mapped
+
+
+def mapped_field(field) -> Any:
+    return field
+
+
 class InnerDoc(ObjectBase, metaclass=DocumentMeta):
     """
     Common class for inner documents like Object or Nested

elasticsearch_dsl/field.py

@@ -77,8 +77,6 @@ def __init__(self, multi=False, required=False, *args, **kwargs):
         """
         self._multi = multi
         self._required = required
-        self._name = None
-        self._parent = None
         super().__init__(*args, **kwargs)
 
     def __getitem__(self, subfield):

examples/async/vectors.py

@@ -47,6 +47,8 @@
 import asyncio
 import json
 import os
+from datetime import datetime
+from typing import List, Optional
 from urllib.request import urlopen
 
 import nltk
@@ -55,12 +57,10 @@
 
 from elasticsearch_dsl import (
     AsyncDocument,
-    Date,
     DenseVector,
     InnerDoc,
     Keyword,
-    Nested,
-    Text,
+    M,
     async_connections,
 )
 
@@ -72,22 +72,22 @@
 
 
 class Passage(InnerDoc):
-    content = Text()
-    embedding = DenseVector()
+    content: M[str]
+    embedding: M[DenseVector]
 
 
 class WorkplaceDoc(AsyncDocument):
     class Index:
         name = "workplace_documents"
 
-    name = Text()
-    summary = Text()
-    content = Text()
-    created = Date()
-    updated = Date()
-    url = Keyword()
-    category = Keyword()
-    passages = Nested(Passage)
+    name: M[str]
+    summary: M[str]
+    content: M[str]
+    created: M[datetime]
+    updated: M[Optional[datetime]]
+    url: M[Keyword]
+    category: M[Keyword]
+    passages: M[Optional[List[Passage]]]
 
     _model = None
 

examples/vectors.py

@@ -46,22 +46,15 @@
 import argparse
 import json
 import os
+from datetime import datetime
+from typing import List, Optional
 from urllib.request import urlopen
 
 import nltk
 from sentence_transformers import SentenceTransformer
 from tqdm import tqdm
 
-from elasticsearch_dsl import (
-    Date,
-    DenseVector,
-    Document,
-    InnerDoc,
-    Keyword,
-    Nested,
-    Text,
-    connections,
-)
+from elasticsearch_dsl import DenseVector, Document, InnerDoc, Keyword, M, connections
 
 DATASET_URL = "https://github.com/raw/elastic/elasticsearch-labs/main/datasets/workplace-documents.json"
 MODEL_NAME = "all-MiniLM-L6-v2"
@@ -71,22 +64,22 @@
 
 
 class Passage(InnerDoc):
-    content = Text()
-    embedding = DenseVector()
+    content: M[str]
+    embedding: M[DenseVector]
 
 
 class WorkplaceDoc(Document):
     class Index:
         name = "workplace_documents"
 
-    name = Text()
-    summary = Text()
-    content = Text()
-    created = Date()
-    updated = Date()
-    url = Keyword()
-    category = Keyword()
-    passages = Nested(Passage)
+    name: M[str]
+    summary: M[str]
+    content: M[str]
+    created: M[datetime]
+    updated: M[Optional[datetime]]
+    url: M[Keyword]
+    category: M[Keyword]
+    passages: M[Optional[List[Passage]]]
 
     _model = None