Update trivial inference for Python 3.14 (#37248)

* Update trivial inference for Python 3.14

* correct comment

* Address review coments

* avoid none case being incorrect

* fix docstring

Author: jrmccluskey

sdks/python/apache_beam/typehints/native_type_compatibility.py

@@ -95,7 +95,7 @@ def _get_args(typ):
     A tuple of args.
   """
   try:
-    if typ.__args__ is None:
+    if typ.__args__ is None or not isinstance(typ.__args__, tuple):
       return ()
     return typ.__args__
   except AttributeError:

sdks/python/apache_beam/typehints/native_type_compatibility_test.py

@@ -337,7 +337,7 @@ def test_forward_reference(self):
     self.assertEqual(typehints.Any, convert_to_beam_type('int'))
     self.assertEqual(typehints.Any, convert_to_beam_type('typing.List[int]'))
     self.assertEqual(
-        typehints.List[typehints.Any], convert_to_beam_type(typing.List['int']))
+        typehints.List[typehints.Any], convert_to_beam_type(list['int']))
 
   def test_convert_nested_to_beam_type(self):
     self.assertEqual(typehints.List[typing.Any], typehints.List[typehints.Any])

sdks/python/apache_beam/typehints/opcodes.py

@@ -63,6 +63,11 @@
 else:
   _div_binop_args = frozenset()
 
+if sys.version_info >= (3, 14):
+  _NB_SUBSCR_OPCODE = [op[0] for op in opcode._nb_ops].index('NB_SUBSCR')
+else:
+  _NB_SUBSCR_OPCODE = -1
+
 
 def pop_one(state, unused_arg):
   del state.stack[-1:]
@@ -151,6 +156,9 @@ def get_iter(state, unused_arg):
 
 def symmetric_binary_op(state, arg, is_true_div=None):
   # TODO(robertwb): This may not be entirely correct...
+  # BINARY_SUBSCR was rolled into BINARY_OP in 3.14.
+  if arg == _NB_SUBSCR_OPCODE:
+    return binary_subscr(state, arg)
   b, a = Const.unwrap(state.stack.pop()), Const.unwrap(state.stack.pop())
   if a == b:
     if a is int and b is int and (arg in _div_binop_args or is_true_div):
@@ -206,7 +214,10 @@ def binary_subscr(state, unused_arg):
       out = base._constraint_for_index(index.value)
     except IndexError:
       out = element_type(base)
-  elif index == slice and isinstance(base, typehints.IndexableTypeConstraint):
+  elif (index == slice or getattr(index, 'type', None) == slice) and isinstance(
+      base, typehints.IndexableTypeConstraint):
+    # The slice is treated as a const in 3.14, using this instead of
+    # BINARY_SLICE
     out = base
   else:
     out = element_type(base)
@@ -483,20 +494,29 @@ def load_global(state, arg):
   state.stack.append(state.get_global(arg))
 
 
+def load_small_int(state, arg):
+  state.stack.append(Const(arg))
+
+
 store_map = pop_two
 
 
 def load_fast(state, arg):
   state.stack.append(state.vars[arg])
 
 
+load_fast_borrow = load_fast
+
+
 def load_fast_load_fast(state, arg):
   arg1 = arg >> 4
   arg2 = arg & 15
   state.stack.append(state.vars[arg1])
   state.stack.append(state.vars[arg2])
 
 
+load_fast_borrow_load_fast_borrow = load_fast_load_fast
+
 load_fast_check = load_fast
 
 
@@ -605,6 +625,8 @@ def set_function_attribute(state, arg):
                     for t in state.stack[attr].tuple_types)
   new_func = types.FunctionType(
       func.code, func.globals, name=func.name, closure=closure)
+  if arg & 0x10:
+    new_func.__annotate__ = attr
   state.stack.append(Const(new_func))
 
 

sdks/python/apache_beam/typehints/trivial_inference.py

@@ -396,6 +396,11 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
 
   jump_multiplier = 2
 
+  # Python 3.14+ push nulls are used to signal kwargs for CALL_FUNCTION_EX
+  # so there must be a little extra bookkeeping even if we don't care about
+  # the nulls themselves.
+  last_op_push_null = 0
+
   last_pc = -1
   last_real_opname = opname = None
   while pc < end:  # pylint: disable=too-many-nested-blocks
@@ -441,7 +446,8 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
         elif op in dis.haslocal:
           # Args to double-fast opcodes are bit manipulated, correct the arg
           # for printing + avoid the out-of-index
-          if dis.opname[op] == 'LOAD_FAST_LOAD_FAST':
+          if dis.opname[op] == 'LOAD_FAST_LOAD_FAST' or dis.opname[
+              op] == "LOAD_FAST_BORROW_LOAD_FAST_BORROW":
             print(
                 '(' + co.co_varnames[arg >> 4] + ', ' +
                 co.co_varnames[arg & 15] + ')',
@@ -450,6 +456,8 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
             print('(' + co.co_varnames[arg & 15] + ')', end=' ')
           elif dis.opname[op] == 'STORE_FAST_STORE_FAST':
             pass
+          elif dis.opname[op] == 'LOAD_DEREF':
+            pass
           else:
             print('(' + co.co_varnames[arg] + ')', end=' ')
         elif op in dis.hascompare:
@@ -512,6 +520,12 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
         # stack[-has_kwargs]: Map of keyword args.
         # stack[-1 - has_kwargs]: Iterable of positional args.
         # stack[-2 - has_kwargs]: Function to call.
+        if arg is None:
+          # CALL_FUNCTION_EX does not take an arg in 3.14, instead the
+          # signaling for kwargs is done via a PUSH_NULL instruction
+          # right before CALL_FUNCTION_EX. A PUSH_NULL indicates that
+          # there are no kwargs.
+          arg = ~last_op_push_null
         has_kwargs: int = arg & 1
         pop_count = has_kwargs + 2
         if has_kwargs:
@@ -680,6 +694,9 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
       jmp_state = state.copy()
       jmp_state.stack.pop()
       state.stack.append(element_type(state.stack[-1]))
+    elif opname == 'POP_ITER':
+      # Introduced in 3.14.
+      state.stack.pop()
     elif opname == 'COPY_FREE_VARS':
       # Helps with calling closures, but since we aren't executing
       # them we can treat this as a no-op
@@ -694,6 +711,10 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
       # We're treating this as a no-op to avoid having to check
       # for extra None values on the stack when we extract return
       # values
+      last_op_push_null = 1
+      pass
+    elif opname == 'NOT_TAKEN':
+      # NOT_TAKEN is a no-op introduced in 3.14.
       pass
     elif opname == 'PRECALL':
       # PRECALL is a no-op.
@@ -727,6 +748,10 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
     else:
       raise TypeInferenceError('unable to handle %s' % opname)
 
+    # Clear check for previous push_null.
+    if opname != 'PUSH_NULL' and last_op_push_null == 1:
+      last_op_push_null = 0
+
     if jmp is not None:
       # TODO(robertwb): Is this guaranteed to converge?
       new_state = states[jmp] | jmp_state