fix(narrow): narrow TypedDict unions on subscript membership tests

[mikeleppane]

Summary

What

Narrow discriminated TypedDict unions when the discriminant is tested with a membership check (event[key] in (...) / not in (...)), not just with equality/identity.

Given a union where each member pins a shared key to a distinct literal:

TraceEvent = PauseEvent | CounterEvent | ProcessMeta | InstantEvent  # ph: "X" | "C" | "M" | "I"

if event["ph"] in ("X", "C", "I"):
    timed.append(event)   # event now narrows to PauseEvent | CounterEvent | InstantEvent

ProcessMeta (ph: Literal["M"]) is now correctly dropped inside the branch, and the union is fully restored in the else.

Why

Subscript narrowing does two jobs: it narrows the facet value (event["ph"]), and it narrows the parent by eliminating union members whose discriminant can't match. The parent step (atomic_narrow_for_facet) handled ==/!=/is/is not but left in/not in to narrow nothing. So equality discriminants worked while membership discriminants silently left the full union in place - impossible members survived the branch, and downstream assignments (list.append, returns, etc.) failed to type-check even though the code is sound.

How

The in/not in arm now narrows each union member's discriminant facet through the same atomic_narrow routine already used for the facet value, and drops any member whose facet collapses to Never - a member that can't satisfy the test is impossible in that branch.

Fixes #3603

Test Plan

cargo test -p pyrefly test::narrow - full narrowing module green
python3 test.py --no-test --no-conformance --no-jsonschema - cargo fmt + clippy clean

[github-actions]

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

bokeh (https://github.com/bokeh/bokeh)
- ERROR src/bokeh/models/layouts.py:332:20-31: Returned type `Nullable[int]` is not assignable to declared return type `int | None` [bad-return]

pydantic (https://github.com/pydantic/pydantic)
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `AfterValidatorFunctionSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `AnySchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ArgumentsSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ArgumentsV3Schema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `BeforeValidatorFunctionSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `BoolSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `BytesSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `CallSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `CallableSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ChainSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ComplexSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `CustomErrorSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DataclassArgsSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DataclassSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DateSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DatetimeSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DecimalSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DefinitionReferenceSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DefinitionsSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `DictSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `EnumSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `FloatSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `FrozenSetSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `GeneratorSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `IntSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `InvalidSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `IsInstanceSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `IsSubclassSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `JsonOrPythonSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `JsonSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `LaxOrStrictSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ListSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `LiteralSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `MissingSentinelSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ModelFieldsSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `ModelSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `NoneSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `NullableSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `PlainValidatorFunctionSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `SetSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `StringSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `TaggedUnionSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `TimeSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `TimedeltaSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `TupleSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `TypedDictSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `UnionSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `UuidSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `WithDefaultSchema` [unsupported-operation]
- ERROR pydantic/networks.py:122:13-45: Cannot set item in `WrapValidatorFunctionSchema` [unsupported-operation]

[github-actions]

Primer Diff Classification

✅ 2 improvement(s) | 2 project(s) total | -2 errors

2 improvement(s) across bokeh, pydantic.

Project	Verdict	Changes	Error Kinds	Root Cause
bokeh	✅ Improvement	-1	bad-return	pyrefly/lib/alt/narrow.rs
pydantic	✅ Improvement	-1	unsupported-operation	atomic_narrow_for_facet()

Detailed analysis

✅ Improvement (2)

bokeh (-1)

This is a removal of a false positive. The _sphinx_height_hint method declares return type int | None, and self.height is declared as Nullable(NonNegative(Int)), which is a Bokeh property descriptor. When accessed on an instance, this descriptor should yield a value of type int | None. Pyrefly was previously reporting the type as Nullable[int] (the descriptor wrapper type) rather than resolving it to the underlying value type int | None, causing a spurious incompatibility error with the declared return type int | None. The PR fixed how pyrefly resolves the value type of property descriptors like Nullable, correctly treating self.height as int | None rather than the descriptor class Nullable[int].
Attribution: The PR changes to pyrefly/lib/alt/narrow.rs added support for in/not in membership narrowing in the AtomicNarrowOp::In(v) | AtomicNarrowOp::NotIn(v) arms. The condition on line 331 self.sizing_mode in ("stretch_width", "fixed", None) is an in membership test. The improved narrowing logic (particularly the non_literal_membership guard and the literal_membership_exprs check) likely changed how the type of self.height is resolved within the narrowed branch, fixing the false positive where Nullable[int] was not being recognized as assignable to int | None.

pydantic (-1)

This is a clear improvement. The PR adds support for narrowing discriminated TypedDict unions on in/not in membership tests. The removed error was a false positive — the code correctly guards with not in ('url', 'multi-host-url') before the subscript assignment, ensuring only URL-related schema variants reach line 122. Pyrefly now correctly narrows the union and no longer reports a spurious unsupported-operation error.
Attribution: The addition of AtomicNarrowOp::In(v) | AtomicNarrowOp::NotIn(v) handling in atomic_narrow_for_facet() in pyrefly/lib/alt/narrow.rs enables TypedDict union narrowing via membership tests. This directly removes the false positive at line 122 where the not in check on line 117 now correctly narrows the schema type.

---

Was this helpful? React with 👍 or 👎

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