Uploaded image for project: 'Beam'
  1. Beam
  2. BEAM-4091

Typehint annotations don't work with @ptransform_fn annotation

    XMLWordPrintableJSON

    Details

    • Type: Bug
    • Status: Resolved
    • Priority: P2
    • Resolution: Fixed
    • Affects Version/s: 2.4.0
    • Fix Version/s: None
    • Component/s: sdk-py-core
    • Labels:

      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)
      

        Attachments

          Issue Links

            Activity

              People

              • Assignee:
                udim Udi Meiri
                Reporter:
                chuanyu Chuan Yu Foo
              • Votes:
                0 Vote for this issue
                Watchers:
                5 Start watching this issue

                Dates

                • Created:
                  Updated:
                  Resolved:

                  Time Tracking

                  Estimated:
                  Original Estimate - Not Specified
                  Not Specified
                  Remaining:
                  Remaining Estimate - 0h
                  0h
                  Logged:
                  Time Spent - 7h 50m
                  7h 50m