No keys hints for TypedDict when it is in union

[pukhov]

Environment data

• Language Server version: Pylance language server 2022.5.3 (pyright 73c82fab)
• OS and version: Windows 10 / tested on Linux, too.
• Python version (& distribution if applicable, e.g. Anaconda): 3.10.4 (anaconda)

Code Snippet

from __future__ import annotations
from typing import TypedDict

class MyDict(TypedDict):
    a: str

def foo() -> MyDict | int:
    ...

x = foo()
# x. # OK: shows hints for int, e.g. bit_length and for dict, e.g. keys
# x[] # BUG: does not show 'a' here
# x[''] # and even here

# but this works, though, so it understands it is TypedDict, only hints do not work.
#y = x['a'] # inferred type of y is str | Any

Expected behavior

Show keys hints, 'a' in this case

Actual behavior

Does not show key hints.

[bschnurr]

I think you need to check that the type is MyDict first

if isinstance(x, MyDict):
    x['a'] = 'hi'
else:
    x = 1

[erictraut]

For type correctness, you'd need to add that check. For completions, we should offer suggestions for all subtypes in the union. We do this for unions in other cases, like in this example.

The completion logic for TypedDict is special-cased, and it doesn't currently handle unions properly.

[pukhov]

@bschnurr

I don't think one should check with isinstance to get autocomplete hints, finally, it's just a nice feature, the inference of types works correctly (I said so in the issue, that y = x['a'] has correct type -- str | Any).

[pukhov]

@erictraut

So it is currently not supported but maybe it will be fixed later?

[erictraut]

It's a trivial fix. Working on it now.

[erictraut]

This will be addressed in the next release.

The change also supports unions that include multiple TypedDict types.

[pukhov]

Thanks!

In fact, what I want to achieve with this is to say that some class has these and these keys. Are there any other ways to tell this to pylance?

As I understand, now only TypedDict has supported hints. There are 2 "pythonic" solutions, both are prohibited by TypedDict pep though :)

1. class should inherit from TypedDict
2. add methods to TypedDict

So what I am doing is a hack, I say "this is both class and it has these keys" by passing the union of my class and created TypedDict.

Motivation / common use-case: pandas.DataFrame / numpy.ndarray (with fields) with known in advance columns / dtype.

I might create a separate issue for this, though not sure, if it's pylance or current Python typing state. (I think the latter)

[erictraut]

This would be a topic better discussed in the python/typing discussions forum.

[pukhov]

Ok, thanks!

[erictraut]

We'll leave this open until next week's release.

[pukhov]

@erictraut

BTW, is the problem below also fixed by your change or should I open another issue?

Normal classes have hints in notebooks, but TypedDicts don't.

foo.py

from typing import TypedDict

class A:
    a: int

class B(TypedDict):
    a: int

def get_A() -> A:
    ...

def get_B() -> B:
    ...

in notebook:
(on double brackets see #2384 )

import foo

a = foo.get_A()
b = foo.get_B()

a. # gives hints
b[] # no hints
b[[]] # and even this

[erictraut]

That should have worked prior to my change because b doesn't involve a union.

Here are the completions I receive with the latest code (which includes my recent change). I think this is what one would expect.

Edit: It looks like there's another issue that may be specific to Jupyter notebooks. There's already another bug tracking that issue.

[pukhov]

Yes, it works perfectly in normal py files, however for notebooks:

• if dict is already created and has keys then there is a known issue, when double brackets work and single brackets do not.
• if it is not created (but of a known type TypedDict), then neither single nor double brackets work, so I thought it's a different issue.

We can wait for that issue to be resolved first and if it still does not work, I ll open a new issue.

[bschnurr]

This issue has been fixed in version 2022.6.0, which we've just released. You can find the changelog here: CHANGELOG.md