1

I have beam code like:


def foo(i: int) -> Foo:
  ...
class IntToFoo(beam.PTransform):
  def expand(self, pcoll: PCollection[int]) -> PCollection[Foo]:
    return pcoll | beam.Map(foo)

Running pyright it complains for the returning line:

Expression of type "PValue" cannot be assigned to return type "PCollection[Foo]"
  "PValue" is incompatible with "PCollection[Foo]

Now AFAIU it seems like beam does not provide sufficient type hinting. Is there a way of working around that as a user of beam so that pyright accepts the code (without ignoring the type error).

EDIT:

I now tried:

    temp = pcoll | beam.Map(foo)
    assert isinstance(temp, PCollection[Foo])
    return temp

Good news is, neither temp = nor return line has any errors. Bad news is that it complains on isinstance line now:

TypeVar or generic type with type arguments not allowed

EDIT 2:

For now I settled with

temp = pcoll | beam.Map(foo)
return cast(PCollection[Foo], temp)

which IMO is still not a solution, because I then would need to specify PCollection[Foo] both the signature as well as the cast.

abergmeier
  • 13,224
  • 13
  • 64
  • 120
  • What happens if you put the isinstance line inside an `if TYPE_CHECKING:` block? (from typing import TYPE_CHECKING). I'm not saying it's a solution to your original question, just curious. – Udi Meiri Mar 16 '21 at 00:27
  • 1
    BTW, is pyright or mypy failing throwing this error? – Udi Meiri Mar 16 '21 at 01:06
  • pyright is complaining - have not tried mypy yet – abergmeier Mar 16 '21 at 12:14
  • fwiw - we have mypy tests running on the Beam codebase, and they work fine. We have not tried pyright on Beam code : / – Pablo Mar 16 '21 at 16:12
  • FYI @Pablo - I created https://issues.apache.org/jira/browse/BEAM-12004 for pyright. – abergmeier Mar 17 '21 at 07:30
  • @Pablo looking at ignores in https://github.com/apache/beam/blob/master/sdks/python/mypy.ini, "mypy tests" might be a bit of a stretch *cough* – abergmeier Mar 18 '21 at 11:33

0 Answers0