[BEAM-4091] Typehint annotations don't work with @ptransform_fn annotation

• Type: Bug
• Status: Triage Needed
• Priority: P2
• Resolution: Fixed
• Affects Version/s: 2.4.0
• Fix Version/s: Missing
• Component/s: sdk-py-core
• Labels: types
• Assignee: Udi Meiri
• Reporter: Chuan Yu Foo
• Created: 16/Apr/18 20:32
• Updated: 13/Apr/23 11:02
• Resolved: 19/Mar/21 17:20

Description

Typehint annotations don't work with functions annotated with @ptransform_fn, but they do work with the equivalent classes.

The following is a minimal example illustrating this:

@beam.typehints.with_input_types(float)
@beam.typehints.with_output_types(bytes)
@beam.ptransform_fn
def _DoStuffFn(pcoll):
  return pcoll | 'TimesTwo' >> beam.Map(lambda x: x * 2)

@beam.typehints.with_input_types(float)
@beam.typehints.with_output_types(bytes)
class _DoStuffClass(beam.PTransform):

  def expand(self, pcoll):
    return pcoll | 'TimesTwo' >> beam.Map(lambda x: x * 2)

With definitions as above, the class correctly fails the typecheck:

def class_correctly_fails():
  p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True))
  _ = (p
       | 'Create' >> beam.Create([1, 2, 3, 4, 5])
       | 'DoStuff1' >> _DoStuffClass()
       | 'DoStuff2' >> _DoStuffClass()
       | 'Write' >> beam.io.WriteToText('/tmp/output'))
  p.run().wait_until_finish()

# apache_beam.typehints.decorators.TypeCheckError: Input type hint violation at DoStuff1: expected <type 'float'>, got <type 'int'>

But the ptransform_fn incorrectly passes the typecheck:

def ptransform_incorrectly_passes():
  p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True))
  _ = (p
       | 'Create' >> beam.Create([1, 2, 3, 4, 5])
       | 'DoStuff1' >> _DoStuffFn()
       | 'DoStuff2' >> _DoStuffFn()
       | 'Write' >> beam.io.WriteToText('/tmp/output'))
  p.run().wait_until_finish()
# No error

Note that changing the order of the @ptransform_fn and type hint annotations doesn't change the result, i.e. changing _DoStuffFn to the following still results in it incorrectly passing the typecheck:

@beam.ptransform_fn
@beam.typehints.with_input_types(float)
@beam.typehints.with_output_types(bytes)
def _DoStuffFn(pcoll):
  return pcoll | 'TimesTwo' >> beam.Map(lambda x: x * 2)

Issue Links

• is duplicated by: BEAM-8502 Python typehints: support type hint decorators with ptransform_fn
• links to: GitHub Pull Request #9907
• links to: GitHub Pull Request #13246