Narrow types based on collection containment

[Noghpu]

Describe the Bug

I just ran into python/mypy#20602 and noticed that pyrefly also doesn't type narrow in some of these cases.

from typing import Literal, reveal_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 no_narrowing(key: str):
    if key in lit_dict:
        reveal_type(key)  # str
        lit_dict[key]  # Argument `str` is not assignable to parameter `key` with type `Literal['a']`
    if key in lit_tuple:
        reveal_type(key)  # str
    if key in lit_list:
        reveal_type(key)  # str
    if key in lit_set:
        reveal_type(key)  # str
    if key in lit_frozenset:
        reveal_type(key)  # str
    if key in frozenset(("a",)):
        reveal_type(key)  # str

def narrowing(key: str):
    if key in ("a",):
        reveal_type(key)  # Literal['a']
    if key in ["a"]:
        reveal_type(key)  # Literal['a']
    if key in {"a"}:
        reveal_type(key)  # Literal['a']
    if key in {"a": 1}:
        reveal_type(key)  # Literal['a']

https://pyrefly.org/sandbox/?project=N4IgZglgNgpgziAXKOBDAdgEwEYHsAeAdAA4CeS4ATrgLYAEALqcROgOZ0Q3G6UN0AZCAxiVUUADR1KMAG4xxAfSbEYAHXQahDAIJ0AvIOGjxAbTUhUFgLoaowxZggBjBojpPXp7TqmsG1gZ0wBZWSHQAjAC%2Bdg4MAK7EsO4JSTDewjqBhgAUoRYSAJSxDIr2cG505QwZutl05pY2JYpwMJVtNT71IU0gMej2pWDUAF4w6J3uI7jjk%2B21WUEzc505eX1FxZo7mDBgdOi4iuiolNQA7qxsOQDWMKTuFZSFiBp0H5wH96Sc6FUOTxud6fUEyeRKFQwO4PQofADEdGeINBHyGjhcNR%2BgToiJ0lDY8RoE34AANnqTOHBDrh%2BKg4HAIGxTthYIxcHRiGdUMSRJQ6KSfpSrgwABaMZgwAXaExQUwAclQ8uspJRXzoPz%2BANKqWSarBcgUUGUkphpDhuKRDEoaog3weWvR1Te-1RH3BRpNqjNFsRyNdHztGodrG1rXaLrd7sNkNNP19VptAfVmtD6JWEym%2Bs%2BHtj3vjCMTtvtv1DGfmDHW%2BRAW0jUdzxqhPsL-o0Gj2B1O51wV3YZqe1texeDpf%2BGzCRTrbobXuhBctMrEcsVyuHqf%2BjTC1inqJnTfniMXZhXtmTQfXwWrUR3BohjbjsMLR%2BXStPoPPIf%2BvTC7miN5zMb3vmj4LsYS4Kq%2BNYgPEDDQHAJDkIgICIgAqrBQy-GA8ToK4EC4JMbZYPsdBgLwNCoKU6BEtgog5Pg7j%2BHCAC0AB8iY7jICSUP8YAWAActRojuMA%2BAxCAGhQWQMhgFApCEAwtBQBQiIAAqkNJslIhgOAEHQzj4ZAhJiLB%2BGEBoiIAMowFKooMAwxBwIgAD0TlSfssmELwbBORMTmYLgzhwE5%2BnoIZ8TGXh6BOSRvB0KgsioNAqCslKIVhRF%2BF0LgxAmZMZnoGQYr4Ux8iUIymWGBYADMhARAATBYElRBIli4fIABi0AwBQaBYHgRBkP0QA

So it seems to narrow when checking against inline collections, but not if we check against typed variables.

Sandbox Link

No response

(Only applicable for extension issues) IDE Information

No response

[yangdanny97]

Some of these cases might be solved by #2627, which is merged and will be available in the next release. There might be some followup items here to handle all the cases you listed

[yangdanny97]

The sandbox is updated w/ the latest changes. Your lit_dict example works as expected now, but the other ones are not supported yet. It shouldn't be too hard to add, building on top of the dict support added in #2627

    if key in lit_tuple:
        reveal_type(key)  # str
    if key in lit_list:
        reveal_type(key)  # str
    if key in lit_set:
        reveal_type(key)  # str
    if key in lit_frozenset:
        reveal_type(key)  # str
    if key in frozenset(("a",)):
        reveal_type(key)  # str

[JanEricNitschke]

Also seeing this for narrowing of "str | None" against a "str" collection:

pyrefly

pyright