fix Preliminary support for class decorators #2752

[asukaminato0721]

Summary

Fixes #2752

applying class decorators to the runtime class-value bindingm while preserving the original class object when a decorator still returns a usable type form.

Test Plan

add test

[Copilot]

Pull request overview

This PR enhances Pyrefly’s solver to type-check class decorator application so the decorated class name can be rebound to the decorator’s return type (e.g., @remote turning a class name into an ActorHandle[...]), and adds a regression test covering this behavior.

Changes:

• Apply non-dataclass-transform class decorators during Binding::ClassDef type solving to infer the post-decoration value type.
• Extend untype_opt/expr_untype handling for Type::KwCall and name lookups to better support the new decorator behavior.
• Add a new decorator test case ensuring a class decorator can rebind the class value.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
pyrefly/lib/test/decorators.rs	Adds a regression testcase for a class decorator that rebinds the class symbol to an ActorHandle[T].
pyrefly/lib/alt/solve.rs	Implements class-decorator application during solving and adjusts untyping/name handling to support the new behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[Copilot]

untype_opt is being used here purely as a predicate to detect whether the decorator result is a type form, but untype_opt calls canonicalize_all_class_types, which can emit user-facing errors (e.g., implicit-Any errors for tuple/Callable). That can lead to spurious errors at class decoration sites when a decorator returns something that happens to be untypeable, even though we’re not actually interpreting the result as an annotation. Consider running this check with an ErrorStyle::Never collector (via error_swallower()), or adding a side-effect-free helper for “is this a type form?” so the predicate doesn’t report errors.

Suggested change

	if self.untype_opt(ty.clone(), range, errors).is_some() {
	let mut swallower = error_swallower();
	if self
	.untype_opt(ty.clone(), range, &mut swallower)
	.is_some()
	{

[yangdanny97]

Overall: This is a clear regression. The PR's class decorator support mishandles decorators whose return types are Unknown/unresolvable, causing 179 false positive errors across the spack project. The decorator @lang.lazy_lexicographic_ordering returns the class itself at runtime, but pyrefly now types the class as ((realf: Unknown) -> Unknown) | Unknown, cascading into missing-attribute, not-a-type, invalid-inheritance, unexpected-keyword, and other errors. None of these are flagged by mypy or pyright.

@migeed-z why does this get classified as "improvement" lol

[NathanTempest]

Thank you for your work on this but wouldn't real-world decorators (@deprecated, @wraps, @lazy_lexicographic_ordering, etc.) have return types that resolve to Unknown or generic callables, so the original class type must be preserved. This is what is causing the regressions.

Tagging @samwgoldman / @stroxler because it is their domain.

This may also be worth scoping down to only support decorators whose return type resolves to a concrete class/type-form, instead of unknowns which could be picked in a follow-up.

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

jax (https://github.com/google/jax)
+ ERROR jax/experimental/array_serialization/serialization_test.py:820:1-821:38: `Unknown` is not assignable to upper bound `type[Any]` of type variable `Typ` [bad-specialization]
+ ERROR jax/experimental/array_serialization/serialization_test.py:950:5-951:38: `Unknown` is not assignable to upper bound `type[Any]` of type variable `Typ` [bad-specialization]
+ ERROR jax/experimental/array_serialization/serialization_test.py:963:5-35: `UserPytreeAPITest.test_custom_node_registration.P` is not assignable to upper bound `Hashable` of type variable `H` [bad-specialization]
+ ERROR jax/experimental/array_serialization/serialization_test.py:968:5-969:43: `Unknown` is not assignable to upper bound `type[Any]` of type variable `Typ` [bad-specialization]

spark (https://github.com/apache/spark)
+ ERROR python/pyspark/testing/tests/test_skip_class.py:20:1-15: Argument `type[SkipClassTests]` is not assignable to parameter `reason` with type `str` in function `unittest.case.skip` [bad-argument-type]

[github-actions]

Primer Diff Classification

❌ 1 regression(s) | ✅ 1 improvement(s) | 2 project(s) total | +5 errors

1 regression(s) across jax. error kinds: Unknown inference failure on class decorators, False Hashable constraint on dataclass. 1 improvement(s) across spark.

Project	Verdict	Changes	Error Kinds	Root Cause
jax	❌ Regression	+4	Unknown inference failure on class decorators	pyrefly/lib/alt/solve.rs
spark	✅ Improvement	+1	class decorator argument type mismatch	pyrefly/lib/alt/solve.rs

Detailed analysis

❌ Regression (1)

jax (+4)

❌ Unknown inference failure on class decorators: 3 errors (lines 820, 950, 968) report Unknown is not assignable to upper bound type[Any] when pyrefly tries to analyze @partial(jax.tree_util.register_dataclass, ...) decorators. Pyrefly fails to resolve the return type of the partial-applied decorator, gets Unknown, then incorrectly flags it. These are inference failures, not real bugs.
❌ False Hashable constraint on dataclass: 1 error (line 963) claims local class P (a @dataclass with a: int = 2) is not assignable to Hashable. Dataclasses generate __hash__ by default (when eq=True and frozen is not explicitly set to False with unsafe_hash=False), so P IS hashable. This is a false positive from the new decorator analysis.

Overall: All 4 new errors are introduced by the PR's new class decorator analysis. The PR added logic to analyze class decorators (previously they were skipped with a TODO comment). However, the new analysis fails to properly resolve types for JAX's register_dataclass and register_static decorators:

• 3 errors show Unknown is not assignable to upper bound type[Any] — this is a classic inference failure pattern where pyrefly cannot resolve the decorator's return type and falls back to Unknown, then incorrectly flags the Unknown as violating a type variable bound.
• 1 error claims a @dataclass class P (with field a: int = 2) is not Hashable. This is technically correct: @dataclass with default settings (eq=True, frozen=False) sets __hash__ = None, making instances unhashable. However, jax.tree_util.register_static is designed to accept such classes, and the Hashable bound on its type variable may be overly strict or may be intended to apply differently. Neither mypy nor pyright flags this, suggesting the JAX type stubs or their handling differs.

All errors are pyrefly-only (neither mypy nor pyright flags them), and all are regressions from the new decorator handling code.

Attribution: The PR changes in pyrefly/lib/alt/solve.rs added class decorator analysis logic in the Binding::ClassDef branch. Previously, class decorators were ignored (TODO: analyze the class decorators). Now pyrefly attempts to call-infer through decorators. The new code calls as_call_target_or_error and call_infer on each decorator, which triggers type checking of the decorator application. For JAX's register_dataclass (which is called via partial(jax.tree_util.register_dataclass, data_fields=..., meta_fields=...)), pyrefly cannot properly resolve the return type, resulting in Unknown types that then fail the bad-specialization check against type variable bounds like type[Any] and Hashable. Specifically:

1. Lines 820-821 and 950-951 and 968-969: @partial(jax.tree_util.register_dataclass, ...) — pyrefly now tries to infer the decorator's return type but gets Unknown, then reports Unknown is not assignable to upper bound type[Any] of type variable Typ. This is an inference failure — the decorator returns the class itself.

2. Line 963: @jax.tree_util.register_static applied to a @dataclass class P — pyrefly now reports that the local class P is not assignable to Hashable. The register_static function likely has a type parameter bounded by Hashable, and pyrefly's new decorator handling is incorrectly checking this constraint. A @dataclass class with only an int field and eq=True (default) generates __eq__ and __hash__, so it IS hashable.

✅ Improvement (1)

spark (+1)

✅ class decorator argument type mismatch: unittest.skip is typed in typeshed as expecting str, but receives type[SkipClassTests] when used without parentheses. This is a real type error based on declared types, also caught by pyright. At runtime it works because CPython's implementation intentionally handles callable arguments, but typeshed doesn't reflect this dual-use behavior. The PR enables class decorator type checking, which correctly identifies this mismatch.

Overall: This is a genuine type error based on the declared types. unittest.skip is typed in typeshed as def skip(reason: str) -> .... Using @unittest.skip without parentheses passes the class SkipClassTests as the reason parameter, which is type[SkipClassTests], not str. Pyright also flags this. At runtime, CPython's implementation of unittest.skip intentionally checks whether the reason argument is callable, and if so, treats it as the test to skip directly (returning a skipped version of it). This is a deliberate dual-use API design, but the typeshed annotations only reflect the str parameter version and don't capture this overloaded behavior. Pyrefly is now correctly analyzing class decorator applications (previously it skipped them entirely), and this error is a true positive based on the declared type signatures.

Attribution: The change in pyrefly/lib/alt/solve.rs in the Binding::ClassDef branch now applies class decorators by calling call_infer with the class type as an argument. Previously, decorators were ignored (TODO: analyze the class decorators). Now pyrefly actually type-checks the decorator application, which means @unittest.skip applied to a class is checked as skip(SkipClassTests), revealing the type mismatch between type[SkipClassTests] and str.

Suggested fixes

Summary: The new class decorator analysis in solve.rs doesn't handle Unknown/failed inference results gracefully, causing 4 pyrefly-only regressions in jax where decorators like @partial(register_dataclass, ...) and @register_static produce Unknown or trigger false constraint violations instead of preserving the original class type.

1. In the Binding::ClassDef branch in solve.rs (around the new decorator loop, ~line 5020-5050), after calling call_infer() to get decorated_ty, add a guard that checks if decorated_ty is Unknown (or contains Unknown). When the decorator inference fails and returns Unknown, the code should fall back to preserving the original class type (ty = self.heap.mk_class_def(cls.dupe())) rather than proceeding to the callable/untype checks. Specifically, after let decorated_ty = self.call_infer(...), add: if decorated_ty.is_unknown() { ty = self.heap.mk_class_def(cls.dupe()); continue; }. This handles the case where pyrefly cannot resolve the decorator's return type (e.g., partial-applied functions) and should preserve the class identity, matching the behavior of mypy and pyright.

Files: pyrefly/lib/alt/solve.rs
Confidence: high
Affected projects: jax
Fixes: Unknown inference failure on class decorators
3 of the 4 jax errors show 'Unknown is not assignable to upper bound type[Any]', which means call_infer returned Unknown for the @partial(register_dataclass, ...) decorator. The Unknown type then fails both the decorator_result_is_callable check and the untype_opt check, so it gets assigned as the decorated type. When this Unknown type is later used, it triggers bad-specialization errors. By detecting Unknown early and falling back to the original class type, we match the previous behavior (where decorators were skipped entirely) for cases where inference fails. This eliminates 3 of the 4 regression errors.

2. In the Binding::ClassDef branch in solve.rs, the decorator loop should also handle the case where call_infer produces errors during type variable bound checking but the decorator is identity-like (returns the same type as input). Specifically, for the @register_static case, the decorator has signature like def register_static[T: Hashable](cls: type[T]) -> type[T]. When pyrefly checks the Hashable bound on a @DataClass class, it incorrectly determines the class is not Hashable. The root issue is that @DataClass with default settings (eq=True, frozen not set) should NOT set hash = None — per Python docs, when eq=True and frozen is not specified (defaults to False), hash is set to None only if eq is defined. But dataclass DOES define eq, so hash IS set to None. However, neither mypy nor pyright flags this. Consider adding the decorator's bound-check errors to a separate error list that is discarded (or only emitted as warnings) when the decorator is identity-like (returns type[T] where T is the input class). Alternatively, treat this as covered by the Unknown fix above if the partial resolution also affects this path.

Files: pyrefly/lib/alt/solve.rs
Confidence: medium
Affected projects: jax
Fixes: False Hashable constraint on dataclass
The 1 remaining error about Hashable is technically correct per the spec (dataclass with eq=True, frozen=False sets hash=None), but neither mypy nor pyright reports it. This suggests the ecosystem treats this more leniently. Since the decorator is identity-like (register_static returns the class), suppressing bound-check errors during decorator application for identity decorators would match ecosystem behavior. This eliminates the remaining 1 regression error.

---

Was this helpful? React with 👍 or 👎

_{Classification by primer-classifier (2 LLM)}