Improve match statement union narrowing/inference

[bergkvist]

Improve inference/narrowing support for union types in match statements.

Fixes #17549

Before:

var: tuple[int, int] | tuple[str, str]

# TODO: we can infer better here.
match var:
    case (42, a):
        reveal_type(a)  # N: Revealed type is "Union[builtins.int, builtins.str]"
    case ("yes", b):
        reveal_type(b)  # N: Revealed type is "Union[builtins.int, builtins.str]"

After:

var: tuple[int, int] | tuple[str, str]

match var:
    case (42, a):
        reveal_type(a)  # N: Revealed type is "builtins.int"
    case ("yes", b):
        reveal_type(b)  # N: Revealed type is "builtins.str"

Related:

• Support for python 3.10 match statement #10191

[Hnasar]

Nice!
Does your fix address any of the other match statement issues?
https://github.com/python/mypy/issues?q=sort%3Aupdated-desc+is%3Aopen+label%3Atopic-match-statement

[bergkvist]

Nice! Does your fix address any of the other match statement issues? https://github.com/python/mypy/issues?q=sort%3Aupdated-desc+is%3Aopen+label%3Atopic-match-statement

From reading through the issues, I believe it should at least solve these:

• Unable to narrow the type of a tagged union when using a match statement #15190
• Type is not being narrowed properly in match statement #17549

I notice there are a bunch of match exhaustiveness issues with narrowing of the type passed on to the next match branches, which this PR does not fix.

[hjwp]

maybe it would address this too? #16835

[bergkvist]

maybe it would address this too? #16835

I don't think so, as I return the "current type" as the "rest type", which means the remainder/rest type (which is what gets passed on to the next match case) doesn't get narrowed properly yet. I have yet to figure out how to do this.

[github-actions]

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

[smurzin]

Hi @bergkvist @Hnasar are there any blockers for landing this PR?
even though it might not address all match narrowing issues it moves it towards a better place.

[lordmauve]

I don't think this is solving the issue in my use case involving TypedDicts:

$ uvx --from 'mypy @ git+https://github.com/bergkvist/mypy@tobias/match-union-narrowing' mypy
    Updated https://github.com/bergkvist/mypy (92d3c91aaf2757706f9ef923650eb342d19af2fb)
      Built mypy @ git+https://github.com/bergkvist/mypy@92d3c91aaf2757706f9ef923650eb342d19af2fb
Installed 3 packages in 17ms
src/pack.py:150: error: "object" has no attribute "encode"  [attr-defined]
src/pack.py:151: error: "object" has no attribute "encode"  [attr-defined]

I have:

match node:
    case {"type": ("assign" | "cond_assign") as op, "key": key, "value": value}:
        kbs = key.encode()
        vbs = value.encode()
        ...

where reveal_type(node) shows

src/pack.py:148: note: Revealed type is "Union[TypedDict('propsast.AssignmentOperation', {'type': Union[Literal['assign'], Literal['cond_assign']], 'key': builtins.str, 'value': builtins.str, 'lineno': builtins.int}), TypedDict('propsast.IncludeOperation', {'type': Union[Literal['include'], Literal['tryinclude']], 'filename': builtins.str, 'lineno': builtins.int})]"

I tried removing the disjunction in the case statement and having a case branch for each literal value but that didn't help.

[Hnasar]

@bergkvist Sorry for the late reply, but can you rebase this on the latest mypy, then we can ping one of the real maintainers to get someone with commit privileges to take a look.