'in' can narrow TypedDict unions

[ikonst]

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']

[ikonst]

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

[ikonst]

@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.

[ikonst]

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.

[ikonst]

@sobolevn can this be merged as is?

[ikonst]

@hauntsaninja ping

[hauntsaninja]

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

[hauntsaninja]

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

[ikonst]

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

[ikonst]

520df60

[hauntsaninja]

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])

[ikonst]

520df60

[hauntsaninja]

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

[ikonst]

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?

[ikonst]

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

[hauntsaninja]

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.

[ikonst]

Gotcha, done: fb564eb

[hauntsaninja]

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

[ikonst]

520df60

[github-actions]

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]

[alicederyn]

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!)

[ikonst]

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

[hauntsaninja]

Thanks for improving this! This is a nice feature