Annotate that a class have these keys available

[pukhov]

I'll first write about my motivation and then move on to question/proposal.

Motivation

Now there is no way to annotate pandas.DataFrame and numpy.ndarray (with fields; so-called structured arrays) with known in advance columns / dtype, while it simplifies both coding and documentation, especially if you have many columns.

As I understand, now only TypedDict has supported hints of type __getitem__, so what I am using is

class MyColumns(TypedDict):
    a: np.ndarray
    b: np.ndarray

MyType = MyColumns | np.ndarray
x: MyType = np.array([(1, 2), (3, 4)], dtype=[('a', 'i8'), ('b', 'i8')])

This is, obviously, a hack, and does not work correctly (because x[:5] should be of type MyType again).

Proposed solution

While I don't like the solution below (and it does not work), as it's pretty cumbersome to type this everytime, that would solve the issue (and kind of ok to put it in .pyi file).

class MyType(np.ndarray):
    a: np.ndarray
    b: np.ndarray

    # many more inputs that should go to this overload, like list[slc, ...], etc
    @typing.overload
    def __getitem__(self, key: slice | int) -> MyType: ...

    @typing.overload
    def __getitem__(self, key: str) -> np.ndarray: ...

Now, when you create x: MyType you can slice it many times, but it will always be of type MyType and show a and b attributes.

Questions

1. Is there a way to annotate columns nicer in general? I searched but did not find it.
2. Regarding second solution: now, as I know, I cannot annotate that 'a' and 'b' are keys, not attributes. It would be nice to have syntax like this, e.g.

class MyType(np.ndarray):
    'a': np.ndarray
    'b': np.ndarray
...

maybe this one in particular won't work as it thinks that 'a' is a type, not an available key.