Support Annotated typing hint

Author: miguelgrinberg

elasticsearch/dsl/document_base.py

@@ -33,6 +33,7 @@
     get_args,
     overload,
 )
+from typing_extensions import _AnnotatedAlias
 
 try:
     from types import UnionType
@@ -343,6 +344,10 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
                 # the field has a type annotation, so next we try to figure out
                 # what field type we can use
                 type_ = annotations[name]
+                type_metadata = []
+                if isinstance(type_, _AnnotatedAlias):
+                    type_metadata = type_.__metadata__
+                    type_ = type_.__origin__
                 skip = False
                 required = True
                 multi = False
@@ -389,6 +394,13 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
                     # use best field type for the type hint provided
                     field, field_kwargs = self.type_annotation_map[type_]  # type: ignore[assignment]
 
+                if name not in attrs:
+                    # if this field does not have a right-hand value, we look in the metadata
+                    # of the annotation to see if we find it there
+                    for md in type_metadata:
+                        if isinstance(md, (_FieldMetadataDict, Field)):
+                            attrs[name] = md
+
                 if field:
                     field_kwargs = {
                         "multi": multi,
@@ -401,7 +413,7 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
                 # this field has a right-side value, which can be field
                 # instance on its own or wrapped with mapped_field()
                 attr_value = attrs[name]
-                if isinstance(attr_value, dict):
+                if isinstance(attr_value, _FieldMetadataDict):
                     # the mapped_field() wrapper function was used so we need
                     # to look for the field instance and also record any
                     # dataclass-style defaults
@@ -490,6 +502,12 @@ def __delete__(self, instance: Any) -> None: ...
 M = Mapped
 
 
+class _FieldMetadataDict(dict):
+    """This class is used to identify metadata returned by the `mapped_field()` function."""
+
+    pass
+
+
 def mapped_field(
     field: Optional[Field] = None,
     *,
@@ -514,13 +532,13 @@ def mapped_field(
     when one isn't provided explicitly. Only one of ``factory`` and
     ``default_factory`` can be used.
     """
-    return {
-        "_field": field,
-        "init": init,
-        "default": default,
-        "default_factory": default_factory,
+    return _FieldMetadataDict(
+        _field=field,
+        init=init,
+        default=default,
+        default_factory=default_factory,
         **kwargs,
-    }
+    )
 
 
 @dataclass_transform(field_specifiers=(mapped_field,))

test_elasticsearch/test_dsl/_async/test_document.py

@@ -27,7 +27,7 @@
 import sys
 from datetime import datetime
 from hashlib import md5
-from typing import Any, ClassVar, Dict, List, Optional
+from typing import Annotated, Any, ClassVar, Dict, List, Optional
 
 import pytest
 from pytest import raises
@@ -530,7 +530,7 @@ def test_document_inheritance() -> None:
     } == MySubDoc._doc_type.mapping.to_dict()
 
 
-def test_child_class_can_override_parent() -> None:
+def test_childdoc_class_can_override_parent() -> None:
     class A(AsyncDocument):
         o = field.Object(dynamic=False, properties={"a": field.Text()})
 
@@ -679,117 +679,139 @@ class TypedDoc(AsyncDocument):
         i1: ClassVar
         i2: ClassVar[int]
 
-    props = TypedDoc._doc_type.mapping.to_dict()["properties"]
-    assert props == {
-        "st": {"type": "text"},
-        "dt": {"type": "date"},
-        "li": {"type": "integer"},
-        "ob": {
-            "type": "object",
-            "properties": {
-                "st": {"type": "text"},
-                "dt": {"type": "date"},
-                "li": {"type": "integer"},
+    class TypedDocAnnotated(AsyncDocument):
+        st: Annotated[str, "foo"]
+        dt: Annotated[Optional[datetime], "bar"]
+        li: Annotated[List[int], "baz"]
+        ob: Annotated[TypedInnerDoc, "qux"]
+        ns: Annotated[List[TypedInnerDoc], "quux"]
+        ip: Annotated[Optional[str], field.Ip()]
+        k1: Annotated[str, field.Keyword(required=True)]
+        k2: Annotated[M[str], field.Keyword()]
+        k3: Annotated[str, mapped_field(field.Keyword(), default="foo")]
+        k4: Annotated[M[Optional[str]], mapped_field(field.Keyword())]  # type: ignore[misc]
+        s1: Annotated[Secret, SecretField()]
+        s2: Annotated[M[Secret], SecretField()]
+        s3: Annotated[Secret, mapped_field(SecretField())]  # type: ignore[misc]
+        s4: Annotated[
+            M[Optional[Secret]],
+            mapped_field(SecretField(), default_factory=lambda: "foo"),
+        ]
+        i1: Annotated[ClassVar, "classvar"]
+        i2: Annotated[ClassVar[int], "classvar"]
+
+    for doc_class in [TypedDoc, TypedDocAnnotated]:
+        props = doc_class._doc_type.mapping.to_dict()["properties"]
+        assert props == {
+            "st": {"type": "text"},
+            "dt": {"type": "date"},
+            "li": {"type": "integer"},
+            "ob": {
+                "type": "object",
+                "properties": {
+                    "st": {"type": "text"},
+                    "dt": {"type": "date"},
+                    "li": {"type": "integer"},
+                },
             },
-        },
-        "ns": {
-            "type": "nested",
-            "properties": {
-                "st": {"type": "text"},
-                "dt": {"type": "date"},
-                "li": {"type": "integer"},
+            "ns": {
+                "type": "nested",
+                "properties": {
+                    "st": {"type": "text"},
+                    "dt": {"type": "date"},
+                    "li": {"type": "integer"},
+                },
             },
-        },
-        "ip": {"type": "ip"},
-        "k1": {"type": "keyword"},
-        "k2": {"type": "keyword"},
-        "k3": {"type": "keyword"},
-        "k4": {"type": "keyword"},
-        "s1": {"type": "text"},
-        "s2": {"type": "text"},
-        "s3": {"type": "text"},
-        "s4": {"type": "text"},
-    }
+            "ip": {"type": "ip"},
+            "k1": {"type": "keyword"},
+            "k2": {"type": "keyword"},
+            "k3": {"type": "keyword"},
+            "k4": {"type": "keyword"},
+            "s1": {"type": "text"},
+            "s2": {"type": "text"},
+            "s3": {"type": "text"},
+            "s4": {"type": "text"},
+        }
 
-    TypedDoc.i1 = "foo"
-    TypedDoc.i2 = 123
+        doc_class.i1 = "foo"
+        doc_class.i2 = 123
+
+        doc = doc_class()
+        assert doc.k3 == "foo"
+        assert doc.s4 == "foo"
+        with raises(ValidationException) as exc_info:
+            doc.full_clean()
+        assert set(exc_info.value.args[0].keys()) == {
+            "st",
+            "k1",
+            "k2",
+            "ob",
+            "s1",
+            "s2",
+            "s3",
+        }
 
-    doc = TypedDoc()
-    assert doc.k3 == "foo"
-    assert doc.s4 == "foo"
-    with raises(ValidationException) as exc_info:
+        assert doc_class.i1 == "foo"
+        assert doc_class.i2 == 123
+
+        doc.st = "s"
+        doc.li = [1, 2, 3]
+        doc.k1 = "k1"
+        doc.k2 = "k2"
+        doc.ob.st = "s"
+        doc.ob.li = [1]
+        doc.s1 = "s1"
+        doc.s2 = "s2"
+        doc.s3 = "s3"
         doc.full_clean()
-    assert set(exc_info.value.args[0].keys()) == {
-        "st",
-        "k1",
-        "k2",
-        "ob",
-        "s1",
-        "s2",
-        "s3",
-    }
 
-    assert TypedDoc.i1 == "foo"
-    assert TypedDoc.i2 == 123
-
-    doc.st = "s"
-    doc.li = [1, 2, 3]
-    doc.k1 = "k1"
-    doc.k2 = "k2"
-    doc.ob.st = "s"
-    doc.ob.li = [1]
-    doc.s1 = "s1"
-    doc.s2 = "s2"
-    doc.s3 = "s3"
-    doc.full_clean()
+        doc.ob = TypedInnerDoc(li=[1])
+        with raises(ValidationException) as exc_info:
+            doc.full_clean()
+        assert set(exc_info.value.args[0].keys()) == {"ob"}
+        assert set(exc_info.value.args[0]["ob"][0].args[0].keys()) == {"st"}
 
-    doc.ob = TypedInnerDoc(li=[1])
-    with raises(ValidationException) as exc_info:
-        doc.full_clean()
-    assert set(exc_info.value.args[0].keys()) == {"ob"}
-    assert set(exc_info.value.args[0]["ob"][0].args[0].keys()) == {"st"}
+        doc.ob.st = "s"
+        doc.ns.append(TypedInnerDoc(li=[1, 2]))
+        with raises(ValidationException) as exc_info:
+            doc.full_clean()
 
-    doc.ob.st = "s"
-    doc.ns.append(TypedInnerDoc(li=[1, 2]))
-    with raises(ValidationException) as exc_info:
+        doc.ns[0].st = "s"
         doc.full_clean()
 
-    doc.ns[0].st = "s"
-    doc.full_clean()
-
-    doc.ip = "1.2.3.4"
-    n = datetime.now()
-    doc.dt = n
-    assert doc.to_dict() == {
-        "st": "s",
-        "li": [1, 2, 3],
-        "dt": n,
-        "ob": {
+        doc.ip = "1.2.3.4"
+        n = datetime.now()
+        doc.dt = n
+        assert doc.to_dict() == {
             "st": "s",
-            "li": [1],
-        },
-        "ns": [
-            {
+            "li": [1, 2, 3],
+            "dt": n,
+            "ob": {
                 "st": "s",
-                "li": [1, 2],
-            }
-        ],
-        "ip": "1.2.3.4",
-        "k1": "k1",
-        "k2": "k2",
-        "k3": "foo",
-        "s1": "s1",
-        "s2": "s2",
-        "s3": "s3",
-        "s4": "foo",
-    }
+                "li": [1],
+            },
+            "ns": [
+                {
+                    "st": "s",
+                    "li": [1, 2],
+                }
+            ],
+            "ip": "1.2.3.4",
+            "k1": "k1",
+            "k2": "k2",
+            "k3": "foo",
+            "s1": "s1",
+            "s2": "s2",
+            "s3": "s3",
+            "s4": "foo",
+        }
 
-    s = TypedDoc.search().sort(TypedDoc.st, -TypedDoc.dt, +TypedDoc.ob.st)
-    s.aggs.bucket("terms_agg", "terms", field=TypedDoc.k1)
-    assert s.to_dict() == {
-        "aggs": {"terms_agg": {"terms": {"field": "k1"}}},
-        "sort": ["st", {"dt": {"order": "desc"}}, "ob.st"],
-    }
+        s = doc_class.search().sort(doc_class.st, -doc_class.dt, +doc_class.ob.st)
+        s.aggs.bucket("terms_agg", "terms", field=doc_class.k1)
+        assert s.to_dict() == {
+            "aggs": {"terms_agg": {"terms": {"field": "k1"}}},
+            "sort": ["st", {"dt": {"order": "desc"}}, "ob.st"],
+        }
 
 
 @pytest.mark.skipif(sys.version_info < (3, 10), reason="requires Python 3.10")