This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 77c2dacd1c1 More precise binary operation inference. (#34305)
77c2dacd1c1 is described below
commit 77c2dacd1c1e3797feb38cbee09c93c275f3067e
Author: Robert Bradshaw <[email protected]>
AuthorDate: Wed Mar 19 07:23:29 2025 -0700
More precise binary operation inference. (#34305)
* More precise binary operation inference.
Handle numeric promotion, sequence multiplication, and fix division for
Python 3.11+.
* Fix older python versions.
They have the opcode module, but no _nb_ops.
* Restore fix for 3.9 and 3.10.
* yapf
---
sdks/python/apache_beam/typehints/opcodes.py | 40 +++++++++++++++++-----
.../typehints/trivial_inference_test.py | 16 +++++++++
2 files changed, 48 insertions(+), 8 deletions(-)
diff --git a/sdks/python/apache_beam/typehints/opcodes.py
b/sdks/python/apache_beam/typehints/opcodes.py
index cd78e397f25..e6fdcdba8e4 100644
--- a/sdks/python/apache_beam/typehints/opcodes.py
+++ b/sdks/python/apache_beam/typehints/opcodes.py
@@ -54,6 +54,15 @@ from apache_beam.typehints.typehints import Union
# method on a C-implemented type will do.
_MethodDescriptorType = type(str.upper)
+if sys.version_info >= (3, 11):
+ import opcode
+ _div_binop_args = frozenset([
+ ix for (ix, (argname, _)) in enumerate(opcode._nb_ops)
+ if 'TRUE_DIVIDE' in argname
+ ])
+else:
+ _div_binop_args = frozenset()
+
def pop_one(state, unused_arg):
del state.stack[-1:]
@@ -137,13 +146,32 @@ def get_iter(state, unused_arg):
state.stack.append(Iterable[element_type(state.stack.pop())])
-def symmetric_binary_op(state, unused_arg):
+_NUMERIC_PROMOTION_LADDER = [bool, int, float, complex]
+
+
+def symmetric_binary_op(state, arg, is_true_div=None):
# TODO(robertwb): This may not be entirely correct...
b, a = Const.unwrap(state.stack.pop()), Const.unwrap(state.stack.pop())
if a == b:
- state.stack.append(a)
+ if a is int and b is int and (arg in _div_binop_args or is_true_div):
+ state.stack.append(float)
+ else:
+ state.stack.append(a)
elif type(a) == type(b) and isinstance(a, typehints.SequenceTypeConstraint):
state.stack.append(type(a)(union(element_type(a), element_type(b))))
+ # Technically these next two will be errors for anything but multiplication,
+ # but that's OK.
+ elif a is int and (b in (bytes, str) or
+ isinstance(b, typehints.SequenceTypeConstraint)):
+ state.stack.append(b)
+ elif b is int and (a in (bytes, str) or
+ isinstance(a, typehints.SequenceTypeConstraint)):
+ state.stack.append(a)
+ elif a in _NUMERIC_PROMOTION_LADDER and b in _NUMERIC_PROMOTION_LADDER:
+ state.stack.append(
+ _NUMERIC_PROMOTION_LADDER[max(
+ _NUMERIC_PROMOTION_LADDER.index(a),
+ _NUMERIC_PROMOTION_LADDER.index(b))])
else:
state.stack.append(Any)
@@ -155,12 +183,8 @@ binary_divide = inplace_divide = symmetric_binary_op
binary_floor_divide = inplace_floor_divide = symmetric_binary_op
-def binary_true_divide(state, unused_arg):
- u = union(state.stack.pop(), state.stack.pop)
- if u == int:
- state.stack.append(float)
- else:
- state.stack.append(u)
+def binary_true_divide(state, arg):
+ return symmetric_binary_op(state, arg, True)
inplace_true_divide = binary_true_divide
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py
b/sdks/python/apache_beam/typehints/trivial_inference_test.py
index 9cea6d5794d..c8b59c7ccbf 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py
@@ -244,6 +244,22 @@ class TrivialInferenceTest(unittest.TestCase):
lambda a,
b: a + b, [typehints.List[int], typehints.List[str]])
+ def testBinOpPromotion(self):
+ self.assertReturnType(int, lambda a, b: a + b, [int, bool])
+ self.assertReturnType(float, lambda a, b: a + b, [int, float])
+ self.assertReturnType(complex, lambda a, b: a + b, [int, complex])
+
+ def testBinOpSequenceMul(self):
+ self.assertReturnType(str, lambda a, b: a * b, [int, str])
+ self.assertReturnType(bytes, lambda a, b: a * b, [bytes, int])
+ self.assertReturnType(
+ typehints.List[float], lambda a, b: a * b, [int,
typehints.List[float]])
+
+ def testDiv(self):
+ # We only support Python 3 now.
+ self.assertReturnType(float, lambda a, b: a / b, [int, int])
+ self.assertReturnType(int, lambda a, b: a // b, [int, int])
+
def testCall(self):
f = lambda x, *args: x
self.assertReturnType(
