Is there a `mypy` extension for beam that understands `.expand`?

[cisaacstern]

Over in leap-stc/cmip6-leap-feedstock#40 (comment) @jbusecke ran into the sort of typo that a type checker should catch:

recipe = (
   ...
   OpenURLWithFSSpec(...)  # emits `OpenFileType`s
   | Preprocessor(...)  # takes `xr.Dataset`s
   ...
)

assuming these transforms had been typed as

class OpenURLWithFSSpec(beam.PTransform):
    def expand(pcoll: beam.PCollection[...]) -> beam.PCollection[OpenFileType]:
        ...

class Preprocessor(beam.PTransform):
    def expand(pcoll: beam.PCollection[xr.Dataset]) -> beam.PCollection[xr.Dataset]:
        ...

I would think mypy should be able to understand to infer input and return types from the object's .expand method? Is there an extension that supports this?

[alxmrs]

Have you tried to see if this issue is caught with --type_check_additional?

https://beam.apache.org/documentation/sdks/python-type-safety/

[cisaacstern]

Have you tried to see if this issue is caught with --type_check_additional?

I have not, thanks for the suggestion. Will report back!