[Bug]: FlatMapTuple with certain callables can result in wrong type hints

[hjtran]

What happened?

I'm not exactly sure what the general condition is for this bug, but I have at least one pipeline where using FlatMapTuple with a lambda results in a weird output type hint that includes type.

import argparse
import logging
import re

import apache_beam as beam
from apache_beam.io import ReadFromText
from apache_beam.io import WriteToText
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import SetupOptions


def main(argv=None, save_main_session=True):
  from typing import Tuple  
  def identity(x: Tuple[str, int]) -> Tuple[str, int]:
    return x
  with beam.Pipeline() as p:
    with beam.Pipeline() as p:
      A = (p
                | "Generate input" >> beam.Create([('P1', [2])])
                # This line
                | "Flat" >> beam.FlatMapTuple(lambda k, vs: [(k, v) for v in vs])
                | "Identity" >> beam.Map(identity))


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  main()

This fails with a type hint violation due to the FlatMap method supposedly returning a pcollection with element type Tuple[type, int].

You can fix this by adjusting the flatmap fn to return [(str(k), v) for v in vs], which leads me to believe there's something wrong in the trivial_inference.py that results in this.

Beam Playground Example:
https://play.beam.apache.org/?sdk=python&shared=mSXtb0psDEN

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[liferoad]

cc @jrmccluskey

[hjtran]

Looking at the debug prints, it sorta looks like there might some indexing issue when we try to LOAD_DEREF for k? Maybe has to do with the load_closure opcode where we offset the index for py311? (I'm not very experienced with looking at disassembled python so I could be totally off here)