Python 3.10 Fix jump offsets opcodes

sdks/python/apache_beam/typehints/opcodes.py

@@ -421,6 +421,12 @@ def delete_fast(state, arg):
  state.vars[arg] = Any # really an error
 
 
+# bpo-43683 Adds GEN_START in Python 3.10, but removed in Python 3.11
+# https://github.com/python/cpython/pull/25138
+def gen_start(state, arg):
+ assert len(state.stack) == 0
+
+
 def load_closure(state, arg):
  state.stack.append(state.get_closure(arg))
 

sdks/python/apache_beam/typehints/trivial_inference.py

@@ -376,6 +376,14 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
  inst_size = 2
  opt_arg_size = 0
 
+ # Python 3.10+: bpo-27129 changes jump offsets to use instruction offsets,
+ # not byte offsets. The offsets were halved (16 bits for instructions vs 8
+ # bits for bytes), so we have to double the value of arg.
+ if (sys.version_info.major, sys.version_info.minor) >= (3, 10):
+ jump_multiplier = 2
+ else:
+ jump_multiplier = 1
+
  last_pc = -1
  while pc < end: # pylint: disable=too-many-nested-blocks
  start = pc
@@ -495,23 +503,23 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
  elif opname == 'YIELD_VALUE':
  yields.add(state.stack[-1])
  elif opname == 'JUMP_FORWARD':
- jmp = pc + arg
+ jmp = pc + arg * jump_multiplier
  jmp_state = state
  state = None
  elif opname == 'JUMP_ABSOLUTE':
- jmp = arg
+ jmp = arg * jump_multiplier
  jmp_state = state
  state = None
  elif opname in ('POP_JUMP_IF_TRUE', 'POP_JUMP_IF_FALSE'):
  state.stack.pop()
- jmp = arg
+ jmp = arg * jump_multiplier
  jmp_state = state.copy()
  elif opname in ('JUMP_IF_TRUE_OR_POP', 'JUMP_IF_FALSE_OR_POP'):
- jmp = arg
+ jmp = arg * jump_multiplier
  jmp_state = state.copy()
  state.stack.pop()
  elif opname == 'FOR_ITER':
- jmp = pc + arg
+ jmp = pc + arg * jump_multiplier
  jmp_state = state.copy()
  jmp_state.stack.pop()
  state.stack.append(element_type(state.stack[-1]))

sdks/python/apache_beam/typehints/trivial_inference_test.py

@@ -36,6 +36,14 @@ def assertReturnType(self, expected, f, inputs=(), depth=5):
  expected,
  trivial_inference.infer_return_type(f, inputs, debug=True, depth=depth))
 
+ # The meaning of Jump Offsets in Python 3.10 was changed.
+ # https://github.com/python/cpython/issues/71316
+ # Reported as a bug in Beam https://github.com/apache/beam/issues/21671
+ def testJumpOffsets(self):
+ fn = lambda x: False
+ wrapper = lambda x, *args, **kwargs: [x] if fn(x, *args, **kwargs) else []
+ self.assertReturnType(typehints.List[int], wrapper, [int])
+
  def testBuildListUnpack(self):
  # Lambda uses BUILD_LIST_UNPACK opcode in Python 3.
  self.assertReturnType(