Mypy index type error with `chex.dataclass`

[hylkedonker]

According to the docs, by default a class wrapped with chex.dataclass can be indexed, because the dataclass becomes compatible with collections.abc.Mapping (because mappable_dataclass=True). However, mypy doesn't seem to understand this. For example:

import chex

@chex.dataclass
class Container:
    foo: float

c = Container(foo=1.)
d = c.foo  # OK.
e = c['foo']  # error: Value of type "Container" is not indexable  [index]

Looking at the code, it seems that this is related to methods such as __getitem__ that are added dynamically with setattr which mypy doesn't recognise. Any ideas how to go around this apart for (i) explicitly silencing the error (ii) using a different method of accessing the variables?

Keep up the good work! Hylke