Optimize runtime type checks and state delta encoding

news/+mutable-proxy-fast-path.performance.md

@@ -0,0 +1 @@
+Reading list, dict, and dataclass elements through state vars is about 3.5x faster: proxies are constructed directly, scalar reads skip the wrapping machinery, and the dataclasses-internal check runs once per iteration.

news/+runtime-type-validation.performance.md

@@ -0,0 +1 @@
+Runtime type checks on state var assignments and computed var results now run in pydantic-core, and state deltas are encoded with orjson, cutting per-event overhead for large lists and dicts by an order of magnitude.

packages/reflex-base/news/+runtime-type-validation.performance.md

@@ -0,0 +1 @@
+Add `runtime_isinstance`, a compiled one-level type check backed by pydantic-core, and `format.json_dumps_compact`, an orjson-backed encoder for state deltas; `orjson` is now a dependency.

packages/reflex-base/pyproject.toml

@@ -8,6 +8,7 @@ authors = [{ name = "Khaleel Al-Adhami", email = "khaleel@reflex.dev" }]
 maintainers = [{ name = "Khaleel Al-Adhami", email = "khaleel@reflex.dev" }]
 requires-python = ">=3.10"
 dependencies = [
+  "orjson >=3.11.3,<4",
   "packaging >=24.2,<27",
   "rich >=13,<16",
   "typing_extensions >=4.13.0",

packages/reflex-base/src/reflex_base/utils/format.py

@@ -10,6 +10,7 @@
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any
 
+import orjson
 from rich.markup import escape as escape_markup
 
 from reflex_base import constants
@@ -697,25 +698,46 @@ def format_library_name(library_fullname: str | dict[str, Any]) -> str:
     return lib
 
 
-_serialize: Callable[[Any], Any] | None = None
+if TYPE_CHECKING:
+    from types import ModuleType
 
+_serializers: ModuleType | None = None
 
-def _get_serialize() -> Callable[[Any], Any]:
-    """Get ``serializers.serialize``, importing it on first use.
+
+def _get_serializers() -> ModuleType:
+    """Get the ``serializers`` module, importing it on first use.
 
     The import cannot live at module scope (``serializers`` imports this
     module), and repeating it per call is measurable on the compile path,
-    so the resolved function is cached.
+    so the resolved module is cached.
 
     Returns:
-        The ``serializers.serialize`` callable.
+        The ``reflex_base.utils.serializers`` module.
     """
-    global _serialize
-    if _serialize is None:
+    global _serializers
+    if _serializers is None:
         from reflex_base.utils import serializers
 
-        _serialize = serializers.serialize
-    return _serialize
+        _serializers = serializers
+    return _serializers
+
+
+def _get_serialize() -> Callable[[Any], Any]:
+    """Get ``serializers.serialize``.
+
+    Returns:
+        The ``serializers.serialize`` callable.
+    """
+    return _get_serializers().serialize
+
+
+# Dataclasses and datetimes keep going through the reflex serializers so their
+# output matches ``json_dumps``; orjson's own rendering of both differs.
+_ORJSON_OPTIONS = (
+    orjson.OPT_NON_STR_KEYS
+    | orjson.OPT_PASSTHROUGH_DATACLASS
+    | orjson.OPT_PASSTHROUGH_DATETIME
+)
 
 
 def json_dumps(obj: Any, **kwargs) -> str:
@@ -734,6 +756,39 @@ def json_dumps(obj: Any, **kwargs) -> str:
     return json.dumps(obj, **kwargs)
 
 
+def json_dumps_compact(obj: Any) -> str:
+    """Serialize an object to compact JSON for the wire.
+
+    Produces the same output as ``json_dumps`` with compact separators (reflex
+    serializers handle non-JSON types), encoded by orjson whenever the payload
+    lets it. State deltas and streamed updates go through here.
+
+    Args:
+        obj: The object to be serialized.
+
+    Returns:
+        The JSON string.
+    """
+    serializers = _get_serializers()
+    if not serializers.overrides_native_json_type():
+        try:
+            encoded = orjson.dumps(
+                obj, default=serializers.serialize, option=_ORJSON_OPTIONS
+            )
+        except TypeError:
+            # orjson rejects integers beyond 64 bits, which json accepts.
+            pass
+        else:
+            # orjson collapses NaN and +/-Infinity to null, but the frontend
+            # expects the bare tokens json emits. A null in the output is
+            # either a None or such a float; only then take the slow path.
+            if b"null" not in encoded:
+                return encoded.decode()
+    return json.dumps(
+        obj, ensure_ascii=False, separators=(",", ":"), default=serializers.serialize
+    )
+
+
 def collect_form_dict_names(form_dict: dict[str, Any]) -> dict[str, Any]:
     """Collapse keys with consecutive suffixes into a single list value.
 

packages/reflex-base/src/reflex_base/utils/serializers.py

@@ -126,6 +126,10 @@ def wrapper(fn: SERIALIZED_FUNCTION) -> SERIALIZED_FUNCTION:
         SERIALIZERS[type_] = fn
         get_serializer.cache_clear()
 
+        global _overrides_native_json_type
+        if types.safe_issubclass(type_, _NATIVE_JSON_TYPES):
+            _overrides_native_json_type = True
+
         # Return the function.
         return fn
 
@@ -166,7 +170,10 @@ def serialize(
     # If there is no serializer, return None.
     if serializer is None:
         if dataclasses.is_dataclass(value) and not isinstance(value, type):
-            return {k.name: getattr(value, k.name) for k in dataclasses.fields(value)}
+            return {
+                name: getattr(value, name)
+                for name in _dataclass_field_names(type(value))
+            }
 
         if get_type:
             return None, None
@@ -181,6 +188,37 @@ def serialize(
     return serialized
 
 
+# Bounded like ``is_mutable_type``'s cache: large enough that an app never
+# rescans a class in practice, without retaining dynamically created ones forever.
+@functools.lru_cache(maxsize=1024)
+def _dataclass_field_names(cls: type) -> tuple[str, ...]:
+    """Get the field names of a dataclass, memoized per class.
+
+    Args:
+        cls: The dataclass type.
+
+    Returns:
+        The names of the dataclass fields, in definition order.
+    """
+    return tuple(field.name for field in dataclasses.fields(cls))
+
+
+# Types orjson encodes itself, matching the serializers below; a serializer an
+# app registers for them or a subclass would be bypassed on the wire, so
+# ``json_dumps_compact`` checks this flag.
+_NATIVE_JSON_TYPES = (Enum, UUID)
+_overrides_native_json_type = False
+
+
+def overrides_native_json_type() -> bool:
+    """Whether an app registered a serializer for Enum, UUID, or a subclass.
+
+    Returns:
+        True if such a serializer exists.
+    """
+    return _overrides_native_json_type
+
+
 @functools.lru_cache
 def get_serializer(type_: type) -> Serializer | None:
     """Get the serializer for the type.
@@ -510,3 +548,8 @@ def serialize_image(image: Img) -> str:
                 mime_type = "image/png"
 
         return f"data:{mime_type};base64,{base64_image}"
+
+
+# The built-in Enum and UUID serializers above render exactly as orjson does;
+# only registrations made after this point count as overrides.
+_overrides_native_json_type = False

packages/reflex-base/src/reflex_base/utils/types.py

@@ -30,6 +30,7 @@
     _eval_type,  # pyright: ignore [reportAttributeAccessIssue]
     _GenericAlias,  # pyright: ignore [reportAttributeAccessIssue]
     _SpecialGenericAlias,  # pyright: ignore [reportAttributeAccessIssue]
+    cast,
     get_args,
     is_typeddict,
 )
@@ -1041,6 +1042,187 @@ def _isinstance(
         raise
 
 
+if find_spec("pydantic_core"):
+    from pydantic_core import SchemaValidator, ValidationError, core_schema
+
+    def _runtime_schema(cls: GenericType, nested: int) -> core_schema.CoreSchema | None:
+        """Build a pydantic-core schema equivalent to ``_isinstance`` for a hint.
+
+        Mirrors ``_isinstance(obj, cls, nested=nested, treat_var_as_type=False)``
+        branch for branch; leaves are ``isinstance`` checks so subclass semantics
+        (including ``bool`` as ``int``) are unchanged.
+
+        Args:
+            cls: The type hint to compile.
+            nested: How many container levels the check descends into.
+
+        Returns:
+            The schema, or None when the hint has no exact schema equivalent and
+            the check must run through ``_isinstance`` instead.
+        """
+        if cls is Any:
+            return core_schema.any_schema()
+        if cls is _Var:
+            return core_schema.is_instance_schema(_Var)
+        if cls is None or cls is type(None):
+            return core_schema.none_schema()
+        if isinstance(cls, TypeAliasTypes):
+            return _runtime_schema(resolve_type_alias(cls), nested)
+
+        origin_attr = getattr(cls, "__origin__", None)
+
+        if origin_attr is Union or (
+            origin_attr is None and isinstance(cls, types.UnionType)
+        ):
+            # A Var instance never satisfies a union in _isinstance and a bare Var
+            # member never matches a non-Var value, so the member contributes nothing.
+            choices = [
+                _runtime_schema(arg, nested)
+                for arg in _get_args_cached(cls)
+                if arg is not _Var
+            ]
+            if not choices or None in choices:
+                return None
+            return core_schema.union_schema(
+                cast(
+                    "list[core_schema.CoreSchema | tuple[core_schema.CoreSchema, str]]",
+                    choices,
+                ),
+                mode="left_to_right",
+            )
+
+        if origin_attr is Literal:
+            return core_schema.literal_schema(list(_get_args_cached(cls)))
+
+        origin = origin_attr if origin_attr is not None else _get_origin_cached(cls)
+
+        if origin is None:
+            if is_typeddict(cls):
+                # Key-level validation of a TypedDict has no core schema with the
+                # same shallow semantics; only the element-level dict check does.
+                return None if nested else core_schema.is_instance_schema(dict)
+            if cls is float:
+                return core_schema.is_instance_schema((float, int))
+            # ``object`` would also admit Var instances, which _isinstance rejects.
+            if not isinstance(cls, type) or cls is object or issubclass(cls, _Var):
+                return None
+            return core_schema.is_instance_schema(cls)
+
+        args = _get_args_cached(cls)
+
+        if not args:
+            return (
+                core_schema.is_instance_schema(origin)
+                if isinstance(origin, type)
+                else None
+            )
+
+        if origin is _Var or origin is _Field:
+            return _runtime_schema(args[0], nested)
+
+        if nested > 0:
+            if origin is list:
+                items = _runtime_schema(args[0], nested - 1)
+                return (
+                    None
+                    if items is None
+                    else core_schema.list_schema(items, strict=True)
+                )
+            if origin is tuple:
+                if args[-1] is Ellipsis:
+                    item = _runtime_schema(args[0], nested - 1)
+                    if item is None:
+                        return None
+                    return core_schema.tuple_schema(
+                        [item], variadic_item_index=0, strict=True
+                    )
+                items = [_runtime_schema(arg, nested - 1) for arg in args]
+                if None in items:
+                    return None
+                return core_schema.tuple_schema(
+                    cast("list[core_schema.CoreSchema]", items), strict=True
+                )
+            if safe_issubclass(origin, Mapping):
+                if origin is not dict:
+                    return None
+                keys = _runtime_schema(args[0], nested - 1)
+                values = _runtime_schema(args[1], nested - 1)
+                if keys is None or values is None:
+                    return None
+                return core_schema.dict_schema(keys, values, strict=True)
+            if origin is set:
+                item = _runtime_schema(args[0], nested - 1)
+                return (
+                    None if item is None else core_schema.set_schema(item, strict=True)
+                )
+
+        base = get_base_class(cls)
+        return core_schema.is_instance_schema(base) if isinstance(base, type) else None
+
+    # Compiled validators keyed by type hint; None marks hints that fall back to
+    # ``_isinstance``. Hints are finite (one per annotated field) so this is unbounded.
+    _RUNTIME_VALIDATORS: dict[Any, SchemaValidator | None] = {}
+
+    def _compile_runtime_validator(cls: GenericType) -> SchemaValidator | None:
+        """Compile the validator for a hint, or None when it needs ``_isinstance``.
+
+        Args:
+            cls: The type hint to compile.
+
+        Returns:
+            The compiled validator, or None for hints without a schema equivalent.
+        """
+        if _Var is _Unloaded:
+            _load_var_classes()
+        schema = _runtime_schema(cls, 1)
+        return None if schema is None else SchemaValidator(schema)
+
+    def runtime_isinstance(obj: Any, cls: GenericType) -> bool:
+        """Check a runtime value against a state var annotation, one level deep.
+
+        Equivalent to ``_isinstance(obj, cls, nested=1, treat_var_as_type=False)``
+        but runs the per-element checks in pydantic-core, compiled once per hint.
+        Hints without an exact schema equivalent, and Var instances, take the
+        ``_isinstance`` path.
+
+        Args:
+            obj: The value to check.
+            cls: The declared type of the value.
+
+        Returns:
+            Whether the value matches the declared type.
+        """
+        try:
+            validator = _RUNTIME_VALIDATORS[cls]
+        except KeyError:
+            validator = _RUNTIME_VALIDATORS[cls] = _compile_runtime_validator(cls)
+        if validator is None or isinstance(obj, _Var):
+            return _isinstance(obj, cls, nested=1, treat_var_as_type=False)
+        if type(obj) is not obj.__class__:
+            # A MutableProxy (wrapt) reports the wrapped value's class, which
+            # isinstance honors but pydantic-core's container checks do not.
+            obj = getattr(obj, "__wrapped__", obj)
+        try:
+            validator.validate_python(obj)
+        except ValidationError:
+            return False
+        return True
+
+else:  # pragma: no cover - pydantic is an optional dependency of reflex-base
+
+    def runtime_isinstance(obj: Any, cls: GenericType) -> bool:
+        """Check a runtime value against a state var annotation, one level deep.
+
+        Args:
+            obj: The value to check.
+            cls: The declared type of the value.
+
+        Returns:
+            Whether the value matches the declared type.
+        """
+        return _isinstance(obj, cls, nested=1, treat_var_as_type=False)
+
+
 def is_dataframe(value: type) -> bool:
     """Check if the given value is a dataframe.