feat: mostly support Generic - #272
Draft
nstarman wants to merge 34 commits into
Draft
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends Plum’s dispatch system to better support parameterized generics (stdlib list[int]-style hints and user-defined Generic[T] classes), including improved matching via __orig_class__, a new @plum.generic decorator for inference, and a new two-tier generic dispatch cache.
Changes:
- Add generic-aware matching and caching (including
__orig_class__-aware bearability checks and a two-tier_generic_cache). - Introduce
plum._generichelpers (is_generic_hint,le_generic) and exportplum.generic. - Add extensive tests, docs, and a docs timing generator for generics performance.
Reviewed changes
Copilot reviewed 18 out of 19 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| uv.lock | Adds plum-dispatch to dependency groups (notably docs). |
| pyproject.toml | Adds plum-dispatch to docs dependency group. |
| noxfile.py | Adds a docs session that regenerates generics timing and builds Jupyter Book. |
| src/plum/_type.py | Enhances type-hint resolution and marks parameterized generics unfaithful. |
| src/plum/_signature.py | Switches signature matching to is_bearable_with_orig. |
| src/plum/_resolver.py | Tracks generic origins and adds an arity-1 pre-filtered resolution shortcut. |
| src/plum/_generic.py | Adds generic-hint helpers and the @plum.generic decorator. |
| src/plum/_function.py | Adds two-tier generic cache and __orig_class__-aware cache keying. |
| src/plum/_bear.py | Adds is_bearable_with_orig for __orig_class__-aware matching. |
| src/plum/init.py | Exports generic. |
| tests/test_generic_dispatch.py | New tests for stdlib generics + user-defined generics + caching behavior. |
| tests/test_generic_decorator.py | New tests for @plum.generic inference, slots, dataclasses, and routing. |
| tests/test_cache.py | Updates cache expectations for mixed faithful+generic overloads. |
| tests/benchmark_generics.py | Adds a benchmark script for generic dispatch and __orig_class__ scenarios. |
| docs/types.md | Adds a note pointing to custom generic type dispatch docs. |
| docs/generics.md | New documentation page describing custom generic dispatch patterns and limitations. |
| docs/_toc.yml | Adds the new generics page to the docs TOC. |
| docs/_scripts/time_generics.py | Generates the timing table included in docs/generics.md. |
| docs/_generated/generics_timing.md | Adds generated timing table output. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Coverage Report for CI Build 29856409773Coverage increased (+0.1%) to 99.594%Details
Uncovered Changes
Coverage RegressionsNo coverage regressions found. Coverage Stats
💛 - Coveralls |
nstarman
force-pushed
the
generics
branch
2 times, most recently
from
May 19, 2026 22:54
54b7895 to
0ebf1aa
Compare
Collaborator
Author
It's pretty performant! My naive implementations were ~20x the time. Now it's 1.6 |
| (i for i, m in enumerate(self.methods) if m.signature == signature), | ||
| None, | ||
| ) | ||
| if existing_idx is not None: |
Collaborator
There was a problem hiding this comment.
why are we dropping this? I mean registering should be something done rarely so it's not a huge issue... or is it?
Collaborator
Author
There was a problem hiding this comment.
Since register() is essentially never on the hot path, I've rolled this part back in c8fde4b.
register() now does the full exhaustive scan again and raises if a new signature is equal to more than one existing method, restoring the previous upstream behavior (raised explicitly rather than via assert, so it isn't stripped under -O):
_equal_indices = [i for i, m in enumerate(self.methods) if m.signature == signature]
if len(_equal_indices) > 1:
raise AssertionError(
f"The added method `{method}` is equal to {len(_equal_indices)} "
f"existing methods. This should never happen."
)
existing_idx = _equal_indices[0] if _equal_indices else NoneThe next(...) first-match early-exit is gone.
The other register() change from this PR — updating the derived metadata (is_faithful, is_faithful_for_non_generic, generic_origins, _arity1_methods) incrementally instead of rescanning all methods on every call — is kept. That one is a genuine O(N²)→O(N) win over a full registration and is unaffected by this concern: it only consumes the resulting existing_idx, and the "REPLACE leaves metadata unchanged" reasoning holds regardless of how the matching index is found (equal signatures ⇒ equal types/arity). Also added a test that asserts the duplicate-signature case raises AssertionError.
Collaborator
Author
|
@wesselb this got to be a doozy, in a push one corner of the rug another corner pops up kind of way. |
nstarman
commented
May 20, 2026
| return isinstance(origin, type) and Generic in origin.__mro__ | ||
|
|
||
|
|
||
| def is_bearable_with_orig(value: object, hint: Any, /) -> bool: |
Collaborator
Author
There was a problem hiding this comment.
@leycec is there a more beartype way to do this? I hacked this together to also get the type parameter info.
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
get_origin(Annotated[X, ...]) returns the inner type X on Python 3.10 (the minimum supported version), not Annotated itself, so the Annotated entry in _EXCLUDED_ORIGINS never matched and Annotated hints were misclassified as generic. Detect them explicitly via __metadata__ (PEP 593) before the origin check. Also remove the now-dead Annotated entry from _EXCLUDED_ORIGINS and the corresponding import. Regression tests added.
__orig_class__ is set by Python after subscripted instantiation and by the @Generic decorator. Document this trust assumption with a comment; no logic change.
Previously __le__ ran two full list-comprehension passes when the equality check failed but a subset relationship held, constructing TypeHintWrapper for every type pair twice (4k constructions for k pairs). Replace the two separate passes with a single pre-built list of (wx, wy) tuples that both the equality check and the subset check reuse, halving TypeHintWrapper constructions to 2k. Switch from all([list_comp]) to all(generator) so the checks themselves also short-circuit. Regression test added: asserts exactly 4 TypeHintWrapper constructions for a two-pair Sig(bool, bool) <= Sig(int, int) comparison. Signed-off-by: nstarman <nstarman@users.noreply.github.com>
In _resolve_method_with_cache, when has_generic_signatures is True, type(a) was called twice per argument on the non-generic hot path: 1. Inside issubclass(type(a), o) in the needs_generic generator. 2. Again in tuple(map(type, args)) for the cache key. Hoist types = tuple(map(type, args)) before the generic check and rewrite the generator to iterate over the pre-computed types tuple (issubclass(t, o)), halving type() calls to once per argument. For generic-enabled functions called with non-generic args this saves n_args extra type() calls on every cache miss and every cache hit path that passes through the generic check. Regression test added: shadows type() in plum._function via patch.dict and asserts exactly 1 call for a single-arg dispatch.
Before building any TypeHintWrapper objects, check three O(1) scalar conditions that prove inequality immediately: 1. len(self.types) != len(other.types) 2. one signature has varargs and the other does not 3. self.precedence != other.precedence Previously the method constructed TypeHintWrapper for every type in both signatures and compared the resulting tuples, paying the full wrapping cost even when a simple length or precedence difference made equality impossible. The precedence check was also redundant inside the tuple; it has been moved to the early-exit guard, leaving only the type and varargs wrappers in the final comparison. Regression test added: asserts 0 TypeHintWrapper calls for length, precedence, and varargs-presence mismatches.
Comparable.is_comparable expands to:
self < other or self == other or self > other
Each branch internally calls Signature.__eq__ (which constructs
TypeHintWrapper objects) in addition to __le__. For equal 1-type
signatures this amounts to 6 TypeHintWrapper constructions:
- self.__le__(other) → 2 constructions (from __lt__ via __le__)
- Signature.__eq__ → 2 constructions (from self != other check)
- Signature.__eq__ again → 2 constructions (from 'or self == other')
The Signature override exploits the fact that two signatures are
comparable iff one is a subtype of the other, so __le__ in each
direction is both necessary and sufficient:
return bool(self.__le__(other)) or bool(other.__le__(self))
This short-circuits after the first direction when it returns True,
reducing the equal-signature case from 6 to 2 TypeHintWrapper
constructions and the incomparable case from 6 to 4.
is_comparable is called on the hot path in both _function.py
(AmbiguousLookupError check) and _resolver.py (candidate filtering).
Regression test added: asserts at most 2 TypeHintWrapper constructions
for is_comparable on equal 1-type signatures.
Replace:
for method in [m for m in methods if check(m)]:
with:
for method in methods:
if not check(method):
continue
The list comprehension built an O(k) temporary list (k = number of
matching methods) before entering the processing loop on every call to
_resolve_from. Direct iteration with an early-continue avoids that
allocation entirely.
_resolve_from is called on every cache miss (via resolve() and
resolve_for_type()), so this reduction matters in workloads with many
unique argument-type combinations.
Regression test added: asserts that _resolve_from source contains the
direct-iteration pattern rather than the list-comprehension form.
…eaks Function._instances was a plain list holding a strong reference to every Function ever created. Functions were therefore never garbage-collected for the lifetime of the process, growing memory unboundedly in long-running applications or test suites that create many dispatch functions. Replace with weakref.WeakSet: - _instances: list["Function"] = [] -> weakref.WeakSet["Function"] - Function._instances.append(self) -> Function._instances.add(self) clear_all_cache() and any other code that iterates _instances is unaffected: WeakSet supports plain iteration and automatically skips dead entries. The manual cleanup call plum.Function._instances.pop(-1) in test_defaults is removed because WeakSet reclaims the short-lived f_wrong_default automatically once it goes out of scope after the test. Tests updated: - test_function: assert g in _instances (WeakSet has no __getitem__) - test_function: new GC test verifies Function is reclaimed after del
_arity1_methods buckets are keyed by parameterised-generic origins (e.g.
'list' from 'list[int]'). _can_match_arity1_origin returns False for
annotations that are neither generic hints nor plain types — most notably
typing.Any and Union[X, Y] — so methods carrying those annotations are
excluded from every origin bucket.
When resolve_for_type gathered only bucket-matched methods and passed
them to _resolve_from, resolution could raise NotFoundLookupError even
though self.resolve(target) — which scans all registered methods — would
have found a matching fallback.
Fix: wrap the _resolve_from call in resolve_for_type with a
try/except NotFoundLookupError that falls back to self.resolve(target).
Regression tests added to test_generic_dispatch.py:
- test_plain_any_fallback_with_generic_overload
list[int] + Any overloads; f([1.0, 2.0]) must route to Any,
not raise NotFoundLookupError.
- test_union_fallback_with_generic_overload
list[int] + Union[list, dict] overloads; same scenario with a
Union annotation that also lands in no origin bucket.
…on overloads The _arity1_methods bucket is keyed by parameterised-generic origins. Methods annotated with plain typing.Any or Union[list, dict] are not generic hints, so they appear in no bucket. When the bucket lookup fails, fall back to self.resolve(target) over all registered methods. Regression tests added for both Any and Union fallback cases.
…back The characterization test must call f([]) *before* any non-ambiguous dispatch so the _generic_cache is cold. A prior f([1]) call warms the (list,) cache bucket with the Sequence[int] candidate; beartype then vacuously accepts [] as Sequence[int], silencing the ambiguity before the resolver is ever reached. Signed-off-by: nstarman <nstarman@users.noreply.github.com>
…patible_with_mypyc
… rescan
The original code recomputed generic_origins and _arity1_methods by
iterating self.methods on every register() call. For N registrations
that is O(N²) total work.
Replace with incremental updates:
- REPLACE path (same signature): faithfulness and generic_origins are
unchanged; only swap the method reference in _arity1_methods buckets.
- APPEND path (new signature):
- is_faithful: O(1) — check new method only
- generic_origins: scan new method's types only (O(types per method))
- _arity1_methods: incremental bucket update when already populated;
one-time full rebuild only on first generic registration
The old _sort_most_specific_first used a layer-peeling loop:
- Computed 'layer = [m for m in remaining if not any(o < m for o in remaining)]'
which called Comparable.__lt__ = '__le__ and __ne__', incurring a
Signature.__eq__ (TypeHintWrapper construction) for every comparable pair.
- Called remaining.remove(m) for each extracted method (O(N) scan each time).
New Kahn's algorithm:
- Pre-computation: one pass over N*(N-1)/2 unordered pairs, 2 '__le__' calls
each (le_ij and le_ji). Strict ordering derived as 'le_ij and not le_ji',
so Signature.__eq__ is never called — eliminates O(N^2) TypeHintWrapper
constructions for comparable pairs.
- BFS processes in O(N + E) with no list.remove() calls.
- Safety valve: if Kahn's queue empties early (cyclic __le__, invalid partial
order), remaining methods are appended in original order.
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
…n dataclasses Python's _GenericAlias.__call__ silently swallows FrozenInstanceError when trying to set __orig_class__ via normal attribute assignment, leaving the inferred value (set by the wrapped __init__) in place. As a result, FrozenBox[str](1) dispatched as FrozenBox[int]. Fix: @Generic installs a thin __setattr__ override on frozen dataclasses that allows __orig_class__ to be written via object.__setattr__, so Python's own machinery can overwrite the inferred value with the subscripted alias — the same 'subscripted wins' behaviour as non-frozen classes. Add a canary test (test_cpython_swallows_frozen_instance_error_for_orig_class) that verifies the CPython bug is still present. If CPython is ever fixed the test will fail and the workaround can be removed. Signed-off-by: nstarman <nstarman@users.noreply.github.com>
Previously @Generic only checked hasattr(cls, '__infer_type_parameter__'), so a plain method or staticmethod would pass decoration and then fail silently at instantiation. Now inspect.getattr_static is used to walk the MRO without invoking descriptors, and TypeError is raised immediately if the attribute is absent or is not a classmethod.
When dispatch_multi registers one function for multiple generic signatures (e.g. list[int] and list[str]), the slow-path cache-append guard was comparing by implementation identity (existing_impl is impl). Because both signatures share the same impl object, the second entry was never appended, leaving the bucket with only one hint_tuple. A subsequent call with an empty list would then match the single entry without ever detecting that a second, equally-specific candidate existed — silently returning instead of raising AmbiguousLookupError. Fix: deduplicate by hint_tuple equality instead. The number of distinct hint_tuples is bounded by the number of registered method signatures, so unbounded bucket growth is not a concern. Add a regression test: dispatch_multi registers _impl for both list[int] and list[str]; after priming the cache with [1] and ['a'], calling f([]) must raise AmbiguousLookupError.
Co-authored-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
Clarify section titles and improve explanations regarding type inference with @plum.generic.
…ister register() had been optimized to stop at the first equal-signature match via next(), which silently overwrote that method and dropped upstream's guard that raised AssertionError when a new signature was equal to more than one existing method. register() is not a hot path, so exhaustive detection of a broken "at most one equal signature" invariant matters more than the early exit. Restore the full scan + AssertionError while keeping the incremental metadata speedups (is_faithful, is_faithful_for_non_generic, generic_origins, _arity1_methods) untouched, since those only consume the resulting index. Also clarify the "never faithful" note on parameterised user-defined Generics in _type.py: it does not mean generic overloads poison dispatch caching, and document the generic-vs-non-generic caching split in docs/types.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The topological sort's docstring promised stable original-index order within each layer, but next_queue was built in discovery order — whichever predecessor freed a node last determined its position. Incomparable methods that reached in-degree 0 in the same layer via different predecessors could therefore come out reordered (e.g. [0,1,2,3] -> [0,1,3,2]). No dispatch behaviour depended on this (_resolve_from recomputes the unique most-specific candidate set order-independently, and _generic_cache prepopulation skips buckets with incomparable methods), but the sort makes the ordering deterministic and the docstring honest at negligible cost — this runs only at registration, never on the hot path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The rebase onto upstream/master combined a newer mypy (1.19.0) and beartype (0.22.8) with this branch's switch to beartype.typing.Protocol, leaving _function.py importing Protocol from both typing and beartype.typing. The newer mypy flagged the duplicate (no-redef) and reported the two type: ignore comments as unused. These two protocols (_DispatchFunction, _BoundFunctionProto) are only used as static type annotations, never isinstance-checked, and _BoundFunctionProto exists specifically to give mypy precise typing. beartype.typing.Protocol types as Any to mypy, defeating that purpose, so use plain typing.Protocol here and drop the now-unused ignores. The runtime-checked protocols in _type.py and _generic.py keep beartype.typing.Protocol. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements support for Generics.
AI Disclosure: most ideas were mine. I used claude opus with R/G TDD to implement and refine the code. Then Copilot reviews to catch bugs and continue iterating.
I've spent some time optimizing dispatching on generics, but it's still slower than normal types, so bleeding into this PR are some various performance optimizations I added as I saw them. This should also speed up the "normal" dispatch routes on non-parametrized types!
Remaning Qs:
__method__that can be fast? @leycec this is for you :).