Dual-use decorator return unions alternative heuristic (#2824)

[grievejia]

Summary:

Dual-use decorators often infer a return type that is the union of the final wrapped callable and the leftover decorator factory branch. Preserving that raw union is not useful at decoration sites: every use of the decorated symbol then has to satisfy both branches, even though only the wrapped callable branch can actually represent the final decorated value.

This change normalizes decorator return unions in apply_function_decorator before we finalize the decorated type. When a union member still looks like a decorator factory, meaning it accepts the decoratee as its first argument and itself returns another callable, we treat that member as an artifact of the dual-use control flow and prune it away. If pruning leaves a single unannotated passthrough callable, we preserve the original decoratee type; otherwise we keep the remaining normalized union.

This keeps the fix local to decorator application instead of adding a broad type-level heuristic for all (*args, **kwargs) -> Any callables. That makes the behavior match the motivating dual-use decorator pattern without changing unrelated callable typing.

NOTE: The other alternative approach to this is D96672235

Differential Revision: D97025245

[meta-codesync]

@grievejia has exported this pull request. If you are a Meta employee, you can view the originating Diff in D97025245.

[github-actions]

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

scikit-learn (https://github.com/scikit-learn/scikit-learn)
- ERROR sklearn/utils/tests/test_validation.py:1758:15-16: Expected 1 positional argument, got 3 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1758:18-19: Expected 2 positional arguments, got 3 in function `f1` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1761:15-16: Expected 1 positional argument, got 4 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1761:18-19: Expected 2 positional arguments, got 4 in function `f1` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1768:15-16: Expected 1 positional argument, got 2 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1768:15-16: Expected 1 positional argument, got 2 in function `f2` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1776:15-16: Expected 1 positional argument, got 2 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1776:15-16: Expected argument `b` to be passed by name in function `f3` [unexpected-positional-argument]
+ ERROR sklearn/utils/tests/test_validation.py:1787:15-16: Expected argument `b` to be passed by name in function `f1` [unexpected-positional-argument]
- ERROR sklearn/utils/tests/test_validation.py:1797:12-13: Expected 0 positional arguments, got 3 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1797:18-19: Expected 2 positional arguments, got 3 in function `A1.__init__` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1800:12-13: Expected 0 positional arguments, got 4 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1800:18-19: Expected 2 positional arguments, got 4 in function `A1.__init__` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1808:12-13: Expected 0 positional arguments, got 3 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1808:18-19: Expected 2 positional arguments, got 3 in function `A2.__init__` [bad-argument-count]
- ERROR sklearn/utils/tests/test_validation.py:1811:12-13: Expected 0 positional arguments, got 4 in function `sklearn.utils.validation._inner_deprecate_positional_args` [bad-argument-count]
+ ERROR sklearn/utils/tests/test_validation.py:1811:18-19: Expected 2 positional arguments, got 4 in function `A2.__init__` [bad-argument-count]

pwndbg (https://github.com/pwndbg/pwndbg)
- ERROR pwndbg/commands/cpsr.py:23:1-61: Argument `(((cpsr_value: Unknown | None = None) -> None) -> (cpsr_value: Unknown | None = None) -> None) | ((cpsr_value: Unknown | None = None) -> None)` is not assignable to parameter with type `((cpsr_value: Unknown | None = None) -> None) -> (cpsr_value: Unknown | None = None) -> None` [bad-argument-type]
- ERROR pwndbg/commands/onegadget.py:27:1-63: Argument `(((show_unsat: bool = False, no_unknown: bool = False, verbose: bool = False) -> None) -> (show_unsat: bool = False, no_unknown: bool = False, verbose: bool = False) -> None) | ((show_unsat: bool = False, no_unknown: bool = False, verbose: bool = False) -> None)` is not assignable to parameter with type `((show_unsat: bool = False, no_unknown: bool = False, verbose: bool = False) -> None) -> (show_unsat: bool = False, no_unknown: bool = False, verbose: bool = False) -> None` [bad-argument-type]

[github-actions]

Primer Diff Classification

✅ 2 improvement(s) | 2 project(s) total | +9, -10 errors

2 improvement(s) across scikit-learn, pwndbg.

Project	Verdict	Changes	Error Kinds	Root Cause
scikit-learn	✅ Improvement	+9, -8	bad-argument-count, unexpected-positional-argument	normalize_decorator_return_union()
pwndbg	✅ Improvement	-2	bad-argument-type	apply_function_decorator()

Detailed analysis

✅ Improvement (2)

scikit-learn (+9, -8)

This is an improvement. The PR fixed pyrefly's handling of dual-use decorators like @_deprecate_positional_args. Previously, pyrefly incorrectly inferred that decorated functions only accepted 1 positional argument (seeing the decorator factory signature instead of the wrapped function). Now it correctly sees the actual function signatures and catches real argument count/keyword-only violations that would fail at runtime. All new errors are genuine bugs that pyright also catches.
Attribution: The changes to normalize_decorator_return_union() and normalize_decorator_value_union() in pyrefly/lib/alt/function.rs fixed the dual-use decorator inference. These functions now properly prune decorator factory branches from union returns, allowing pyrefly to see the actual decorated function signatures instead of the decorator factory signature.

pwndbg (-2)

These errors were false positives where pyrefly incorrectly rejected valid dual-use decorators. The OnlyWithArch decorator in pwndbg likely returns a union of the decorator factory and the wrapped function, which confused pyrefly's decorator application logic. The PR implements a heuristic to handle this common pattern by normalizing decorator return unions - when a union contains both a decorator factory branch and a wrapped callable, it keeps only the wrapped callable for the final decorated type. This matches how developers expect dual-use decorators to work and aligns with other type checkers' behavior.
Attribution: The changes to apply_function_decorator() and addition of normalize_decorator_return_union() in pyrefly/lib/alt/function.rs fixed this by pruning decorator factory branches from union returns when applying decorators, allowing the decorated functions to have the correct final type.

---

Was this helpful? React with 👍 or 👎

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