support for optional, list, and other type hints

Author: miguelgrinberg

elasticsearch_dsl/document_base.py

@@ -17,9 +17,21 @@
 
 from datetime import date, datetime
 from fnmatch import fnmatch
+from typing import List, Optional
 
 from .exceptions import ValidationException
-from .field import Field, Integer, Float, Boolean, Text, Binary, Date
+from .field import (
+    Binary,
+    Boolean,
+    Date,
+    Field,
+    Float,
+    InstrumentedField,
+    Integer,
+    Nested,
+    Object,
+    Text,
+)
 from .mapping import Mapping
 from .utils import DOC_META_FIELDS, ObjectBase
 
@@ -35,6 +47,11 @@ def __new__(cls, name, bases, attrs):
         attrs["_doc_type"] = DocumentOptions(name, bases, attrs)
         return super().__new__(cls, name, bases, attrs)
 
+    def __getattr__(cls, attr):
+        if attr in cls._doc_type.mapping:
+            return InstrumentedField(attr, cls._doc_type.mapping[attr])
+        return super().__getattribute__(attr)
+
 
 class DocumentOptions:
     type_annotation_map = {
@@ -53,18 +70,38 @@ def __init__(self, name, bases, attrs):
         # create the mapping instance
         self.mapping = getattr(meta, "mapping", Mapping())
 
-        for name, type_ in attrs.get('__annotations__', {}).items():
-            if name not in attrs:
-                if type_ in self.type_annotation_map:
+        annotations = attrs.get("__annotations__", {})
+        fields = set([n for n in attrs if isinstance(attrs[n], Field)])
+        fields.update(annotations.keys())
+        for name in fields:
+            if name in attrs:
+                value = attrs[name]
+            else:
+                type_ = annotations[name]
+                required = True
+                multi = False
+                while hasattr(type_, "__origin__"):
+                    if type_.__origin__ == Optional:
+                        required = False
+                        type_ = type_.__args__[0]
+                    elif issubclass(type_.__origin__, List):
+                        multi = True
+                        type_ = type_.__args__[0]
+                if issubclass(type_, InnerDoc):
+                    field = Nested if multi else Object
+                    field_args = {}
+                elif type_ in self.type_annotation_map:
                     field, field_args = self.type_annotation_map[type_]
-                    self.mapping.field(name, field(**field_args))
-                elif issubclass(type_, Field):
-                    self.mapping.field(name, type_())
-
-        # register all declared fields into the mapping
-        for name, value in list(attrs.items()):
-            if isinstance(value, Field):
-                self.mapping.field(name, value)
+                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
+            self.mapping.field(name, value)
+            if name in attrs:
                 del attrs[name]
 
         # add all the mappings for meta fields

elasticsearch_dsl/field.py

@@ -77,6 +77,8 @@ 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):
@@ -123,6 +125,25 @@ def to_dict(self):
         return value
 
 
+class InstrumentedField:
+    def __init__(self, name, field):
+        self._name = name
+        self._field = field
+
+    def __getattr__(self, attr):
+        f = None
+        try:
+            f = self._field[attr]
+        except KeyError:
+            pass
+        if isinstance(f, Field):
+            return InstrumentedField(f"{self._name}.{attr}", f)
+        return getattr(self._field, attr)
+
+    def __repr__(self):
+        return self._name
+
+
 class CustomField(Field):
     name = "custom"
     _coerce = True

elasticsearch_dsl/utils.py

@@ -499,6 +499,12 @@ def __getattr__(self, name):
                 return value
             raise
 
+    def __setattr__(self, name, value):
+        if name in self.__class__._doc_type.mapping:
+            self._d_[name] = value
+        else:
+            super().__setattr__(name, value)
+
     def to_dict(self, skip_empty=True):
         out = {}
         for k, v in self._d_.items():