Python: How to define a type annotation so that the resulting object has both attributes and keys with the same name?

Question

I would like to annotate an object obj so that the type checker (or the language server protocol) understand that it has both some attributes and keys with same name, giving me the correct intellisense.

For example:
If obj.foo is a string and obj.bar is an integer: Then, obj["foo"] is a string and obj["bar"] is an integer.

And I would like to get the appropriate intellisense for both attributes/keys when I type obj. or obj["].

I'm using VS Code, with pylance/pyright

Below is what I got so far:

(Obs: the try/except block is needed because there is a run time error when base classes for TypedDict are not TypedDict)

myobj = ... # this object comes from other context

from typing import Protocol, TypedDict, cast

class MyDict(TypedDict):
    foo: str
    bar: int

class MyProtocol(Protocol):
    foo: str
    bar: int

try:
    class MyClass(MyDict, MyProtocol): ...
    myobj = cast(MyClass, myobj)
except TypeError:
    pass

Current use case for this: streamlit session state

If there is a solution for this use case, I would like to know.

Answer 1

You can use multiple @overloaded __getitem__s with Literal:

from typing import Any, Literal, overload

class C:
    foo: int
    bar: bytes
    
    @overload
    def __getitem__(self, n: Literal['foo']) -> int: ...
    
    @overload
    def __getitem__(self, n: Literal['bar']) -> bytes: ...
    
    def __getitem__(self, n: str) -> Any: ...

Screenshots:

Note that this doesn't work if you overload __getattr__ instead.