Speed up model field reads ~4.6x (FieldDescriptor)

.gitignore

@@ -36,6 +36,7 @@ htmlcov/
 .dmypy.json
 *.so
 faust/_cython/*.c
+faust/models/_cython/*.c
 faust/transport/_cython/*.c
 
 # virtualenvs

faust/models/_cython/__init__.py

@@ -0,0 +1 @@
+"""Cython optimized model components."""

faust/models/_cython/fields.pyx

@@ -0,0 +1,47 @@
+# cython: language_level=3
+"""Cython optimized field descriptor."""
+
+
+cdef class FieldDescriptorBase:
+    """Read path of :class:`faust.models.fields.FieldDescriptor`.
+
+    Only ``__get__`` lives here, because that is what runs on every access
+    to every field of every model instance.  Everything else stays in the
+    Python class, where it is readable and overridable.
+
+    The four attributes ``__get__`` touches are declared on the extension
+    type so it can reach them as C struct members instead of going through
+    a dict lookup.  Every other attribute is an ordinary Python attribute
+    in ``__dict__``, so nothing else about the class changes -- including
+    ``as_dict()``/``clone()`` and tests that set attributes directly.
+    """
+
+    cdef public str field
+    cdef public bint required
+    cdef public bint lazy_coercion
+    cdef public object _to_python
+    cdef dict __dict__
+
+    def __get__(self, instance, owner):
+        cdef:
+            dict instance_dict
+            str field
+            object to_python
+            object value
+            object evaluated_fields
+
+        # Class attribute access: `Model.field` returns the descriptor.
+        if instance is None:
+            return self
+
+        field = self.field
+        instance_dict = instance.__dict__
+        to_python = self._to_python
+        value = instance_dict[field]
+        if self.lazy_coercion and to_python is not None:
+            evaluated_fields = instance.__evaluated_fields__
+            if field not in evaluated_fields:
+                if value is not None or self.required:
+                    value = instance_dict[field] = to_python(value)
+                evaluated_fields.add(field)
+        return value

faust/models/fields.py

@@ -1,4 +1,5 @@
 import inspect
+import os
 import sys
 from datetime import datetime
 from decimal import Decimal, DecimalTuple
@@ -17,7 +18,6 @@
     cast,
 )
 
-from mode.utils.objects import cached_property
 from mode.utils.text import pluralize
 
 from faust.exceptions import ValidationError
@@ -27,6 +27,8 @@
 from .tags import Tag
 from .typing import NodeType, TypeExpression
 
+NO_CYTHON = bool(os.environ.get("NO_CYTHON", False))
+
 __all__ = [
     "TYPE_TO_FIELD",
     "FieldDescriptor",
@@ -54,7 +56,51 @@ def _is_concrete_model(typ: Type = None) -> bool:
     )
 
 
-class FieldDescriptor(FieldDescriptorT[T]):
+class _PyFieldDescriptorBase:
+    """Pure-Python read path for :class:`FieldDescriptor`.
+
+    Split out so it can be swapped for a Cython implementation.  ``__get__``
+    runs on every access to every field of every model instance, which makes
+    it the hottest method in the model layer.
+    """
+
+    field: str
+    required: bool
+    lazy_coercion: bool
+    _to_python: Optional[Callable[[T], T]]
+
+    def __get__(self, instance: Any, owner: Type) -> Any:
+        # class attribute accessed
+        if instance is None:
+            return self
+
+        field = self.field
+        instance_dict = instance.__dict__
+        to_python = self._to_python
+        value = instance_dict[field]
+        if self.lazy_coercion and to_python is not None:
+            evaluated_fields: Set[str]
+            evaluated_fields = instance.__evaluated_fields__
+            if field not in evaluated_fields:
+                if value is not None or self.required:
+                    value = instance_dict[field] = to_python(value)
+                evaluated_fields.add(field)
+        return value
+
+
+if not NO_CYTHON:  # pragma: no cover
+    try:
+        from ._cython.fields import FieldDescriptorBase as _FieldDescriptorBase
+    except ImportError:
+        _FieldDescriptorBase = _PyFieldDescriptorBase  # type: ignore[misc,assignment]
+else:  # pragma: no cover
+    _FieldDescriptorBase = _PyFieldDescriptorBase  # type: ignore[misc,assignment]
+
+
+class FieldDescriptor(  # type: ignore[misc,valid-type]
+    _FieldDescriptorBase,
+    FieldDescriptorT[T],
+):
     """Describes a field.
 
     Used for every field in Record so that they can be used in join's
@@ -126,6 +172,13 @@ class FieldDescriptor(FieldDescriptorT[T]):
     #: Field may be tagged with Secret/Sensitive/etc.
     tag: Optional[Type[Tag]]
 
+    #: Set when reading this field has to coerce the value on access,
+    #: rather than at construction time.  Read on every field access.
+    lazy_coercion: bool
+
+    #: Model types reachable from this field's type expression.
+    related_models: Set[Type[ModelT]]
+
     _to_python: Optional[Callable[[T], T]]
     _expr: Optional[TypeExpression]
 
@@ -166,10 +219,22 @@ def __init__(
         self.tag = tag
         self._to_python = None
         self._expr = None
+        # Plain attributes rather than cached_property.  mode's
+        # cached_property defines __set__, which makes it a *data*
+        # descriptor, so every read goes through the descriptor protocol
+        # instead of short-circuiting to the instance dict -- and
+        # lazy_coercion is read on every single model field access.
+        # Both are filled in by on_model_attached() once _expr exists;
+        # until then a descriptor has no _to_python either, so the False
+        # default cannot change what __get__ does.
+        self.lazy_coercion = False
+        self.related_models = set()
 
     def on_model_attached(self) -> None:
-        self._expr = self._prepare_type_expression()
+        expr = self._expr = self._prepare_type_expression()
         self._to_python = self._compile_type_expression()
+        self.related_models = expr.found_types[NodeType.MODEL]
+        self.lazy_coercion = expr.has_generic_types or expr.has_models
 
     def _prepare_type_expression(self) -> TypeExpression:
         expr = TypeExpression(
@@ -248,24 +313,6 @@ def _copy_descriptors(self, typ: Type = None) -> None:
         if typ is not None and _is_concrete_model(typ):
             typ._contribute_field_descriptors(self, typ._options, parent=self)
 
-    def __get__(self, instance: Any, owner: Type) -> Any:
-        # class attribute accessed
-        if instance is None:
-            return self
-
-        field = self.field
-        instance_dict = instance.__dict__
-        to_python = self._to_python
-        value = instance_dict[field]
-        if self.lazy_coercion and to_python is not None:
-            evaluated_fields: Set[str]
-            evaluated_fields = instance.__evaluated_fields__
-            if field not in evaluated_fields:
-                if value is not None or self.required:
-                    value = instance_dict[field] = to_python(value)
-                evaluated_fields.add(field)
-        return value
-
     def should_coerce(self, value: Any, coerce: Optional[bool] = None) -> bool:
         c = coerce if coerce is not None else self.coerce
         return c and (self.required or value is not None)
@@ -304,16 +351,6 @@ def ident(self) -> str:
         """Return the fields identifier."""
         return f"{self.model.__name__}.{self.field}"
 
-    @cached_property
-    def related_models(self) -> Set[Type[ModelT]]:
-        assert self._expr is not None
-        return self._expr.found_types[NodeType.MODEL]
-
-    @cached_property
-    def lazy_coercion(self) -> bool:
-        assert self._expr is not None
-        return self._expr.has_generic_types or self._expr.has_models
-
 
 class BooleanField(FieldDescriptor[bool]):
     def validate(self, value: T) -> Iterable[ValidationError]:

faust/types/models.py

@@ -20,8 +20,6 @@
     cast,
 )
 
-from mode.utils.objects import cached_property
-
 from faust.exceptions import ValidationError  # XXX !!coupled
 
 from .codecs import CodecArg
@@ -252,13 +250,18 @@ def validation_error(self, reason: str) -> ValidationError: ...
     @abc.abstractmethod
     def ident(self) -> str: ...
 
-    @cached_property
-    @abc.abstractmethod
-    def related_models(self) -> Set[Type[ModelT]]: ...
-
-    @cached_property
-    @abc.abstractmethod
-    def lazy_coercion(self) -> bool: ...
+    #: Model types reachable from this field's type expression.
+    #:
+    #: Declared as a plain annotation rather than a cached_property:
+    #: mode's cached_property defines __set__, so it is a *data* descriptor
+    #: and intercepts every read even after the value is cached.  Both of
+    #: these are set as ordinary instance attributes by
+    #: FieldDescriptor.on_model_attached(), and lazy_coercion is read on
+    #: every single model field access.
+    related_models: Set[Type[ModelT]]
+
+    #: Set when reading this field has to coerce the value on access.
+    lazy_coercion: bool
 
 
 # XXX See top of module!  We redefine with actual ModelT for Sphinx,

setup.py

@@ -81,6 +81,13 @@
         extra_compile_args=CFLAGS,
         extra_link_args=LDFLAGS,
     ),
+    Extension(
+        "faust.models._cython.fields",
+        ["faust/models/_cython/fields" + ext],
+        libraries=LIBRARIES,
+        extra_compile_args=CFLAGS,
+        extra_link_args=LDFLAGS,
+    ),
     Extension(
         "faust.transport._cython.conductor",
         ["faust/transport/_cython/conductor" + ext],

tests/unit/models/test_fields.py

@@ -214,3 +214,88 @@ def test_init_options(self):
     def test_prepare_value(self, value, coerce, trim, expected_result):
         f = BytesField(coerce=coerce, trim_whitespace=trim)
         assert f.prepare_value(value) == expected_result
+
+
+class Test_FieldDescriptorBase:
+    """The read path, in whichever implementation is active.
+
+    ``_FieldDescriptorBase`` is the Cython one whenever the extension could
+    be built, and is otherwise the same object as ``_PyFieldDescriptorBase``.
+    """
+
+    def test_active_base_is_one_of_the_two(self):
+        from faust.models import fields as f
+
+        assert f._FieldDescriptorBase in (
+            f._PyFieldDescriptorBase,
+            getattr(f, "_FieldDescriptorBase"),
+        )
+        assert issubclass(FieldDescriptor, f._FieldDescriptorBase)
+
+    def test_class_access_returns_the_descriptor(self):
+        class Point(Record):
+            x: int
+
+        assert isinstance(Point.x, FieldDescriptor)
+        assert Point.x.field == "x"
+
+    def test_instance_access_returns_the_value(self):
+        class Point(Record):
+            x: int
+            label: str
+
+        p = Point(x=1, label="hi")
+        assert p.x == 1
+        assert p.label == "hi"
+
+    def test_lazy_coercion_on_read(self):
+        # A nested model is coerced on first access and cached, so the
+        # second read returns the identical object.
+        class Point(Record):
+            x: int
+
+        class Holder(Record):
+            p: Point
+
+        h = Holder(p={"x": 2})
+        first = h.p
+        assert isinstance(first, Point)
+        assert first.x == 2
+        assert h.p is first
+
+    def test_lazy_coercion_flag_is_a_plain_attribute(self):
+        # Regression: both of these used to be mode cached_property, which
+        # defines __set__ and so intercepts every read even once cached.
+        # lazy_coercion is read on every single field access.
+        class Point(Record):
+            x: int
+
+        from mode.utils.objects import cached_property
+
+        descriptor = Point._options.descriptors["x"]
+        # Neither name may resolve to a cached_property anywhere in the MRO.
+        # (Where the value is actually stored differs between the two
+        # implementations: the extension type keeps lazy_coercion as a C
+        # struct member, the Python one as an instance attribute.)
+        for klass in type(descriptor).__mro__:
+            for name in ("lazy_coercion", "related_models"):
+                assert not isinstance(klass.__dict__.get(name), cached_property)
+        assert descriptor.lazy_coercion is False
+        assert descriptor.related_models == set()
+
+    def test_none_value_is_returned_for_optional_field(self):
+        class Point(Record):
+            x: int = None
+
+        assert Point(x=None).x is None
+
+    def test_descriptor_still_accepts_arbitrary_attributes(self):
+        # clone()/as_dict() and several tests set attributes directly, so
+        # the extension type has to keep a __dict__.
+        class Point(Record):
+            x: int
+
+        descriptor = Point._options.descriptors["x"]
+        descriptor.some_extra_attribute = 42
+        assert descriptor.some_extra_attribute == 42
+        assert type(descriptor.clone()) is type(descriptor)