From 1cbc92aa1c8230710b929d806a248622eddb6477 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sun, 8 Mar 2026 20:13:31 +0900 Subject: [PATCH 1/2] fix containment narrowing now derives finite membership domains directly from typed RHS containers instead of relying only on the generic iterable/mapping decomposition. That preserves literal precision for typed dict, tuple, list, set, and frozenset values, and the literal fast path now also recognizes inline frozenset(...) calls --- pyrefly/lib/test/narrow.rs | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pyrefly/lib/test/narrow.rs b/pyrefly/lib/test/narrow.rs index 67925c7fba..68f028780a 100644 --- a/pyrefly/lib/test/narrow.rs +++ b/pyrefly/lib/test/narrow.rs @@ -2602,6 +2602,35 @@ def non_literal_arg(x: int | str) -> None: "#, ); +testcase!( + test_narrow_in_typed_collections, + r#" +from typing import Literal, assert_type + +LitA = Literal["a"] +lit_dict: dict[LitA, int] = {"a": 1} +lit_tuple: tuple[LitA] = ("a",) +lit_list: list[LitA] = ["a"] +lit_set: set[LitA] = {"a"} +lit_frozenset: frozenset[LitA] = frozenset(("a",)) + +def test(key: str) -> None: + if key in lit_dict: + assert_type(key, LitA) + assert_type(lit_dict[key], int) + if key in lit_tuple: + assert_type(key, LitA) + if key in lit_list: + assert_type(key, LitA) + if key in lit_set: + assert_type(key, LitA) + if key in lit_frozenset: + assert_type(key, LitA) + if key in frozenset(("a",)): + assert_type(key, LitA) +"#, +); + testcase!( test_narrow_len, r#" From 9168f090dd2d11920da2b7e253239f283c95f86c Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 17 Mar 2026 05:03:15 +0900 Subject: [PATCH 2/2] update add more narrow test add remaining case fix TypedDict key narrow --- pyrefly/lib/alt/narrow.rs | 1 + pyrefly/lib/test/narrow.rs | 90 ++++++++++++++++++++++++++++++++++ pyrefly/lib/test/typed_dict.rs | 17 +++++++ 3 files changed, 108 insertions(+) diff --git a/pyrefly/lib/alt/narrow.rs b/pyrefly/lib/alt/narrow.rs index cc085e0bdd..2299540d88 100644 --- a/pyrefly/lib/alt/narrow.rs +++ b/pyrefly/lib/alt/narrow.rs @@ -1570,6 +1570,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // If so, we can narrow the left operand to the mapping's key type. if !self.behaves_like_any(&right_ty) && let Some((key_ty, _)) = self.unwrap_mapping(&right_ty) + && !self.behaves_like_any(&key_ty) { return self.intersect(ty, &key_ty); } diff --git a/pyrefly/lib/test/narrow.rs b/pyrefly/lib/test/narrow.rs index 68f028780a..6198a5c5d9 100644 --- a/pyrefly/lib/test/narrow.rs +++ b/pyrefly/lib/test/narrow.rs @@ -2631,6 +2631,96 @@ def test(key: str) -> None: "#, ); +testcase!( + test_narrow_in_iterable_element_type, + r#" +from typing import assert_type + +def generate_missing() -> list[str]: + return ["a", "b"] + +def takes_string(arg: str) -> None: ... + +def test_list(arg: str | None) -> None: + missing = generate_missing() + if arg in missing: + assert_type(arg, str) + takes_string(arg) + +def test_set(arg: str | None, missing: set[str]) -> None: + if arg in missing: + assert_type(arg, str) + takes_string(arg) + +def test_frozenset(arg: str | None, missing: frozenset[str]) -> None: + if arg in missing: + assert_type(arg, str) + takes_string(arg) +"#, +); + +testcase!( + test_narrow_in_does_not_widen, + r#" +from typing import Any, Literal, Mapping, assert_type + +list_any: list[Any] = [1] +list_object: list[object] = [1] +set_any: set[Any] = {1} +set_object: set[object] = {1} +dict_any: dict[Any, int] = {1: 1} +dict_object: dict[object, int] = {1: 1} +mapping_any: Mapping[Any, int] = {1: 1} +mapping_object: Mapping[object, int] = {1: 1} + +def test_list(x: Literal[1] | str) -> None: + if x in list_any: + assert_type(x, Literal[1] | str) + if x in list_object: + assert_type(x, Literal[1] | str) + +def test_list_optional(x: str | None) -> None: + if x in list_any: + assert_type(x, str | None) + if x in list_object: + assert_type(x, str | None) + +def test_set(x: Literal[1] | str) -> None: + if x in set_any: + assert_type(x, Literal[1] | str) + if x in set_object: + assert_type(x, Literal[1] | str) + +def test_dict(x: Literal[1] | str) -> None: + if x in dict_any: + assert_type(x, Literal[1] | str) + if x in dict_object: + assert_type(x, Literal[1] | str) + +def test_mapping(x: Literal[1] | str) -> None: + if x in mapping_any: + assert_type(x, Literal[1] | str) + if x in mapping_object: + assert_type(x, Literal[1] | str) +"#, +); + +testcase!( + test_in_abstract_mapping_narrows_key, + r#" +from typing import Mapping, assert_type + +def parse_resource_id() -> tuple[str, str] | tuple[None, None]: ... + +def lookup_resource(registry: Mapping[str, str]) -> str | None: + kind, _obj_id = parse_resource_id() + if kind not in registry: + return None + assert_type(kind, str) + return registry[kind] +"#, +); + testcase!( test_narrow_len, r#" diff --git a/pyrefly/lib/test/typed_dict.rs b/pyrefly/lib/test/typed_dict.rs index 1b58cd67c8..78b832540b 100644 --- a/pyrefly/lib/test/typed_dict.rs +++ b/pyrefly/lib/test/typed_dict.rs @@ -2508,6 +2508,23 @@ def test(foo: Foo | Bar) -> None: "#, ); +testcase!( + test_typed_dict_union_key_contains_does_not_narrow, + r#" +from typing import TypedDict, assert_type + +class Foo(TypedDict, closed=True): + a: int + +class Bar(TypedDict, closed=True): + b: int + +def test(schema: Foo | Bar, key: str) -> None: + if key in schema: + assert_type(key, str) +"#, +); + testcase!( test_typed_dict_open_contains_narrowing, r#"