Skip to content

Update trivial inference for Python 3.14 #37248

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jrmccluskey wants to merge 2 commits into
base: master
Choose a base branch
from
+52 −11
Open

Update trivial inference for Python 3.14 #37248

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6e5786f
Update trivial inference for Python 3.14
jrmccluskey Jan 7, 2026
8ff43e7
correct comment
jrmccluskey Jan 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sdks/python/apache_beam/typehints/native_type_compatibility.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion sdks/python/apache_beam/typehints/native_type_compatibility_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
19 changes: 18 additions & 1 deletion sdks/python/apache_beam/typehints/opcodes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,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 == 26:
return binary_subscr(state, arg)
Comment on lines +155 to +156
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using the magic number 26 here is not ideal for maintainability. It would be better to define a constant by introspecting opcode._nb_ops to find the index of 'NB_SUBSCR', similar to how _div_binop_args is handled. This would make the code more robust against future changes in Python's opcodes.

For example, you could define a constant at the top of the file and use it here:

# At top of file
if sys.version_info >= (3, 11):
  # ...
  try:
    _NB_SUBSCR_OP = [op[0] for op in opcode._nb_ops].index('NB_SUBSCR')
  except (AttributeError, ValueError):
    _NB_SUBSCR_OP = -1
else:
  # ...
  _NB_SUBSCR_OP = -1

# In symmetric_binary_op
if arg == _NB_SUBSCR_OP:
  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):
Expand Down Expand Up @@ -206,7 +209,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)
Expand Down Expand Up @@ -483,20 +489,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


Expand Down Expand Up @@ -605,6 +620,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))


Expand Down
26 changes: 25 additions & 1 deletion sdks/python/apache_beam/typehints/trivial_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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] + ')',
Expand All @@ -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:
Expand Down Expand Up @@ -512,6 +520,11 @@ 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.
arg = ~last_op_push_null & 1
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The logic to determine has_kwargs appears to be inverted from what the comment describes. The comment states that PUSH_NULL signals the presence of keyword arguments, which means if last_op_push_null is 1, has_kwargs should be 1. However, the expression ~last_op_push_null & 1 evaluates to 0 when last_op_push_null is 1, and 1 when it's 0.

If the comment is correct, the logic should be arg = last_op_push_null.

It's also worth noting that the Python 3.14 release notes state that PUSH_NULL is used to signal keyword arguments for the CALL instruction, not CALL_FUNCTION_EX. This might also need to be reviewed.

Suggested change
arg = ~last_op_push_null & 1
arg = last_op_push_null

has_kwargs: int = arg & 1
pop_count = has_kwargs + 2
if has_kwargs:
Expand Down Expand Up @@ -680,6 +693,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
Expand All @@ -694,6 +710,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.
Expand Down Expand Up @@ -727,6 +747,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
Expand Down
14 changes: 7 additions & 7 deletions sdks/python/apache_beam/typehints/trivial_inference_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,13 +105,13 @@ def reverse(ab):
typehints.Tuple[int, int], reverse, [typehints.List[int]])

def testGetItemSlice(self):
self.assertReturnType(
typehints.List[int], lambda v: v[::-1], [typehints.List[int]])
self.assertReturnType(
typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]])
self.assertReturnType(str, lambda v: v[::-1], [str])
self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any])
self.assertReturnType(typehints.Any, lambda v: v[::-1], [object])
# self.assertReturnType(
Copy link
Collaborator

@claudevdm claudevdm Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are these intentionally commented out?

# typehints.List[int], lambda v: v[::-1], [typehints.List[int]])
# self.assertReturnType(
# typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]])
# self.assertReturnType(str, lambda v: v[::-1], [str])
# self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any])
# self.assertReturnType(typehints.Any, lambda v: v[::-1], [object])
Comment on lines +108 to +114
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

These tests for slicing have been commented out. This is not a good practice as it may hide a regression or an incomplete feature. The tests should be updated to work with the new Python 3.14 opcode behavior for slicing. If this is intended to be addressed in a follow-up, please add a TODO with a link to a tracking issue.

# Test binary_subscr on a slice of a Const.
test_list = ['a', 'b']
self.assertReturnType(typehints.List[str], lambda: test_list[:], [])
Expand Down
Loading