Avoid copying entire mappings during type inference

packages/reflex-base/news/+bounded-mapping-inference.performance.md

@@ -0,0 +1 @@
+Avoid copying entire mappings when sampling their key and value types during compilation.

packages/reflex-base/src/reflex_base/vars/base.py

@@ -19,6 +19,7 @@
 from collections.abc import Callable, Coroutine, Iterable, Mapping, Sequence
 from dataclasses import _MISSING_TYPE, MISSING
 from decimal import Decimal
+from itertools import islice
 from types import CodeType, FunctionType
 from typing import (
     TYPE_CHECKING,
@@ -1986,8 +1987,8 @@ def figure_out_type(value: Any) -> types.GenericType:
             if not value:
                 return Mapping[NoReturn, NoReturn]
             return Mapping[
-                unionize(*{figure_out_type(k) for k in list(value.keys())[:100]}),
-                unionize(*{figure_out_type(v) for v in list(value.values())[:100]}),
+                unionize(*{figure_out_type(k) for k in islice(value.keys(), 100)}),
+                unionize(*{figure_out_type(v) for v in islice(value.values(), 100)}),
             ]
     return type(value)
 

tests/benchmarks/test_type_inference.py

@@ -0,0 +1,17 @@
+"""Benchmarks for inference from container values."""
+
+import pytest
+from pytest_codspeed import BenchmarkFixture
+from reflex_base.vars.base import figure_out_type
+
+
+@pytest.mark.parametrize("size", [100, 10_000, 100_000])
+def test_mapping_type_inference(benchmark: BenchmarkFixture, size: int):
+    """Infer a mapping's type without work proportional to its size.
+
+    Args:
+        benchmark: The benchmark fixture.
+        size: The number of entries in the mapping.
+    """
+    value = {str(index): index for index in range(size)}
+    benchmark(figure_out_type, value)

tests/units/reflex_base/vars/test_base.py

@@ -2,11 +2,18 @@
 
 import threading
 import typing
+from collections.abc import Mapping
 from typing import Any, Literal, TypeVar
 
 import pytest
 from reflex_base.utils.types import get_field_type
-from reflex_base.vars.base import EvenMoreBasicBaseState, Var, _linearize_bases, field
+from reflex_base.vars.base import (
+    EvenMoreBasicBaseState,
+    Var,
+    _linearize_bases,
+    field,
+    figure_out_type,
+)
 from reflex_base.vars.object import ObjectVar
 from reflex_base.vars.sequence import ArrayVar, StringVar
 from typing_extensions import TypeAliasType, TypeVarTuple, Unpack
@@ -16,6 +23,40 @@
 _MARKER_ATTR = "_marker"
 
 
+@pytest.mark.parametrize("size", [0, 1, 100, 101, 1000])
+def test_mapping_type_inference_bounds_iteration(size):
+    """Sampling does not iterate or retrieve values beyond the first 100 items."""
+    keys_seen = []
+    values_seen = []
+
+    class CountingMapping(Mapping):
+        def __len__(self):
+            return size
+
+        def __iter__(self):
+            for index in range(size):
+                keys_seen.append(index)
+                yield index
+
+        def __getitem__(self, key):
+            values_seen.append(key)
+            return str(key)
+
+    expected = Mapping[int, str] if size else Mapping[typing.NoReturn, typing.NoReturn]
+    assert figure_out_type(CountingMapping()) == expected
+    sampled = list(range(min(size, 100)))
+    assert keys_seen == sampled * 2
+    assert values_seen == sampled
+
+
+def test_mapping_type_inference_preserves_sample_boundary():
+    """The hundredth item contributes types, while later items do not."""
+    value: dict[Any, Any] = dict.fromkeys(range(99), 1)
+    value["last sampled"] = "included"
+    value[1.5] = ["not sampled"]
+    assert figure_out_type(value) == Mapping[int | str, int | str]
+
+
 def test_custom_field_attr_survives_annotated_rebuild():
     """A custom attribute on an annotated Field survives a rebuild."""
     f = field("x")