'in' can narrow TypedDict unions by ikonst · Pull Request #13838 · python/mypy · GitHub

in could narrow unions of TypeDicts, e.g.

class A(TypedDict)
   foo: int

@final
class B(TypedDict):
   bar: int

union: Union[A, B] = ...
value: int
if 'foo' in union:
  # Cannot be a B as it is final and has no "foo" field, so must be an A
  value = union['foo']
else:
  # Cannot be an A as those went to the if branch
  value = union['bar']

Once done™️ , will address #7981 (and #13802).

@hauntsaninja I still need to actually respect @final but am I on the right track?

BTW, does it strictly require @final? This can be deemed as truthy even without @final

class D(TypedDict):
  foo: str

d: D
if 'foo' in d:
  assert_type(d, D)
else:
  assert_never()

It doesn't make a good case for ruling out unions, though:

class D1(TypedDict):
  foo: str

class D2(TypedDict):
  bar: str

d: D1 | D2
if 'foo' in d:
  assert_type(d, Union[D1, D2])
else:
  assert_type(d, D2)

Might seem a bit complex, but we already have to deal with it for totality.

Update: Also discerning on final-ity now

In the case of https://github.com/Rapptz/discord.py/blob/3bca40352ecfa5b219bae0527252ee5d296113e7/discord/client.py#L1755:
it seems that data is typed AppInfo instead of AppInfo | PartialAppInfo, and so we assume rpc_origins is always assigned (true for AppInfo which is total) and deem the branch unreachable.

@sobolevn can this be merged as is?

@hauntsaninja ping

Took a look at the tests and left some comments. I'll review the actual code too, but might take me a little bit to get to it. Will make sure to review before 1.0 though

this test case seems completely duplicated in testOperatorContainsNarrowsTypedDicts_total, let's remove this?

I'll remove the relevant part from testOperatorContainsNarrowsTypedDicts_total, as to keep test names descriptive.

520df60

Suggested change

      
            else:
          
                assert_type(d_or_list, list[str])
          
            elif 'bar' in d_or_list:
          
                assert_type(d_or_list, list[str])
          
            else:
          
                assert_type(d_or_list, list[str])

520df60

in theory this could narrow foo_or_invalid as well, want to add an assert_type for that behaviour too?

That'd have to be implemented. And I think it could be pretty neat, but would require a rework: I'd have to pass in the left expression, and return type maps (not "if_type" and "else_type"). Maybe in a follow-up?

My current implementation is to make a form of "tagged TypeDicts" work, but I suspect a more generalized form of type narrowing should be possible. However, before we tackle the more contrived x in y case, we should probably do it for x == y, so that this example would see narrowing applied.

(I might be naive, though, and this might've been tried before and proven impossible.)

P.S. PyRight is similarly limited in this

Oh sorry, to be clear my suggestion wasn't to implement this / I don't think it's a terribly important feature. I was just saying it's worth adding an assert_type in this test to make this (lack of) behaviour clear to the reader.

Gotcha, done: fb564eb

I think this is because mypy processes these twice. Want to change the test to just do reveal_type(arg['foo'])?

520df60

…nst/mypy into typed-dict-in-type-narrowing

…narrowing

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

discord.py (https://github.com/Rapptz/discord.py)
- discord/client.py:1755: error: Value of "rpc_origins" has incompatible type "None"; expected "List[str]"  [typeddict-item]

Could the PR description be modified to reflect the change made in f799344 ? Perhaps to

class A(TypedDict)
   foo: int

@final
class B(TypedDict):
   bar: int

union: Union[A, B] = ...
value: int
if 'foo' in union:
  # Cannot be a B as it is final and has no "foo" field, so must be an A
  value = union['foo']
else:
  # Cannot be an A as those went to the if branch
  value = union['bar']

(Super excited to see this PR, thank you so much!)

@alicederyn thanks for spotting this — updated with your proposed sample code

Thanks for improving this! This is a nice feature

ikonst force-pushed the typed-dict-in-type-narrowing branch from e714a52 to 4642a9c Compare October 8, 2022 04:03