Preserve TypeVar narrowing across concrete dispatch calls

[mikeleppane]

Summary

What

Pyrefly rejected a dispatch-on-concrete-subclass pattern that pyright accepts. A function over a TypeVar bound to a union, narrowing with isinstance/type(x) is C and returning a helper's result, got a spurious bad-return.

T = TypeVar("T", bound=A | B)
def func_a(a: A) -> A: return a

def dispatch(x: T) -> T:
    if isinstance(x, A):
        return func_a(x)   # before: bad-return (A not assignable to T)
    raise NotImplementedError

Fix in pyrefly/lib/alt/call.rs — refine_typevar_bound_call_result. Plus two scope expansions over the base fix: keyword args + any bounded TypeVar (not just union bounds). Tests in pyrefly/lib/test/narrow.rs.

Why

Type theory: inside the body T is rigid (skolemized), constraint T <: A | B.

• T <: A — using a T as A = fine (down).
• A <: T — false in general: caller's T could be a strict subtype A2 <: A. Can't fabricate a T from a bare A (up = unsound).

Narrowing already does the right thing: isinstance(x, A) gives the meet A & T (reveal_type(x) = A & T), which is both <: A (passable to func_a) and <: T (returnable).

The leak is at the call boundary. func_a: (A) -> A is monomorphic — applying it computes codomain A and drops the & T. So func_a(A & T) typed as plain A → not assignable to -> T → spurious error. Narrowing the variable can't fix it; the value comes from the helper's signature, not from x.

How

After computing a call's result ty, refine it back to the witness — but only when the argument provably is that witness. Gates, in order (cheap → expensive):

1. return-context hint is a single bounded TypeVar (Restriction::Bound(_))
2. exactly one positional or one keyword arg (unified extraction)
3. expected_types[arg_range] == ty — helper's param type == return type (the identity shape; reuses already-computed data, previously discarded as _expected_types)
4. arg's type is an Intersect carrying both that same skolem T and the return type ty → return the intersection (A & T) instead of bare A

A & T <: T → checks. Expensive expr_infer gated behind cheap checks.

Soundness = gates 3+4. A function with param == return behaves (parametrically) like it returns its input; when the input carries & T, the output's precise type still carries & T — we recover the witness, not invent it. Anything else (fresh A, helper returning a subtype (A)->A3) fails a gate → stays rejected.

Honest residual: def f(a: A) -> A: return A() (returns fresh, not input) is accepted though mildly unsound — inherent (the signature can't say "returns its arg"), matches pyright.

Why these two expansions are safe: keyword f(a=x) ≡ positional f(x) (one value, one param) — zero new soundness surface. Bound(Union) → Bound(_) — gates 3+4 don't depend on the bound's shape, so per-instance soundness is identical; only applicability widens. Constrained TypeVars (TypeVar("T", A, B)) already worked soundly via exact-identity and need no special-casing.

Test plan

5 testcases in narrow.rs:

Test	Asserts	Type
test_typevar_dispatch_on_concrete_class	type(x) is A + isinstance dispatch OK; func_a(A()) fresh → bad-return	+/−
test_typevar_dispatch_keyword_arg	func_a(a=x) OK	+
test_typevar_dispatch_single_bound	non-union T: P OK; reveal_type(x) = C1 & Tb	+ (pins shape)
test_typevar_dispatch_negatives	(A)->A3 subtype-return → rejected; func_a(a=A()) fresh keyword → rejected	−

Verification:

• cargo test -p pyrefly 'test::narrow::' → 199 passed, 0 failed
• Full python3 test.py → formatting ✅, lint ✅, 5286 lib tests ✅, conformance ✅. (jsonschema phase couldn't run locally — missing jsonschema/toml pip deps; env gap, not code.)
• mypy_primer ×3 runs, ~17 ecosystem projects (pydantic, pandera, sphinx, aiohttp, starlette, tornado, werkzeug, mypy, attrs, sympy, trio, jinja, packaging, rich…) → zero diagnostic diffs

Fixes #3550

[github-actions]

According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅

[samwgoldman]

Honest residual: def f(a: A) -> A: return A() (returns fresh, not input) is accepted though mildly unsound — inherent (the signature can't say "returns its arg"), matches pyright.

This was my first thought when I saw the initial example. A function A -> A when called with T & A should still return A, since the body could be anything.

I'm not convinced we should follow Pyright's lead on this example.

Edit to add: there is already a way to describe a function that returns its argument, generics. Making func_a generic over U: A ensures that you return the argument by parametricity:

from typing import TypeVar
class A: ...
class B: ...

def func_u[U: A](a: U) -> U:
    return a

def dispatch[T: A | B](x: T) -> T:
    if isinstance(x, A):
        return func_u(x)
    raise NotImplementedError

x = dispatch(B())
reveal_type(x)