From 2f2f1d69003211c62b0c4bed3f8c74a413c46188 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 7 Oct 2026 11:49:15 -0400 Subject: [PATCH 01/13] Add stable identifier pickling option to context --- sdks/python/apache_beam/pipeline.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index fea42bbc0cb3..0e551703266d 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -1081,6 +1081,8 @@ def to_runner_api( # because other runners do not have this option. context.enable_best_effort_deterministic_pickling = getattr( self.runner, 'enable_best_effort_deterministic_pickling', False) + context.enable_stable_identifier_pickling = getattr( + self.runner, 'enable_stable_code_identifier_pickling', False) # The RunnerAPI spec requires certain transforms and side-inputs to have KV # inputs (and corresponding outputs). From 9c03511bcf02f0076698499b44d1cd1b4f11371e Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 7 Oct 2026 11:51:13 -0400 Subject: [PATCH 02/13] Refactor config assignment in cloudpickle_pickler Replace config assignment with dataclasses.replace for stability. --- sdks/python/apache_beam/internal/cloudpickle_pickler.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index cea4f01f803c..86b94fbf3b82 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -29,6 +29,7 @@ import base64 import bz2 +import dataclass import io import logging import sys @@ -187,7 +188,10 @@ def _dumps( with _pickle_lock: with io.BytesIO() as file: if enable_stable_code_identifier_pickling: - config = STABLE_CODE_IDENTIFIER_CONFIG + config = dataclasses.replace( + config, + get_code_object_params=STABLE_CODE_IDENTIFIER_CONFIG + .get_code_object_params) pickler = cloudpickle.CloudPickler(file, config=config) try: pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags From 5477ec1efb3f6361887f31731601637943dafad8 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 7 Oct 2026 13:16:11 -0400 Subject: [PATCH 03/13] Fix import statement for dataclasses module --- sdks/python/apache_beam/internal/cloudpickle_pickler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index 86b94fbf3b82..2d9404feab2f 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -29,7 +29,7 @@ import base64 import bz2 -import dataclass +import dataclasses import io import logging import sys From 88ae832d0d9a336b4dd81f2c7f58571b52c59b37 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 7 Oct 2026 14:30:01 -0400 Subject: [PATCH 04/13] Fix indentation in cloudpickle_pickler.py --- sdks/python/apache_beam/internal/cloudpickle_pickler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index 2d9404feab2f..9710878f5b67 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -189,9 +189,9 @@ def _dumps( with io.BytesIO() as file: if enable_stable_code_identifier_pickling: config = dataclasses.replace( - config, - get_code_object_params=STABLE_CODE_IDENTIFIER_CONFIG - .get_code_object_params) + config, + get_code_object_params=STABLE_CODE_IDENTIFIER_CONFIG. + get_code_object_params) pickler = cloudpickle.CloudPickler(file, config=config) try: pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags From aeeba6c8eeb91697c8a63628ea387c4dc097b72c Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 12:50:36 -0400 Subject: [PATCH 05/13] Revert config changes --- sdks/python/apache_beam/internal/cloudpickle_pickler.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index 9710878f5b67..cea4f01f803c 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -29,7 +29,6 @@ import base64 import bz2 -import dataclasses import io import logging import sys @@ -188,10 +187,7 @@ def _dumps( with _pickle_lock: with io.BytesIO() as file: if enable_stable_code_identifier_pickling: - config = dataclasses.replace( - config, - get_code_object_params=STABLE_CODE_IDENTIFIER_CONFIG. - get_code_object_params) + config = STABLE_CODE_IDENTIFIER_CONFIG pickler = cloudpickle.CloudPickler(file, config=config) try: pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags From 7fd251ca6d025ad5da9e38e9332c92131bf15f15 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 12:51:32 -0400 Subject: [PATCH 06/13] Revert changes --- sdks/python/apache_beam/pipeline.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py index 0e551703266d..fea42bbc0cb3 100644 --- a/sdks/python/apache_beam/pipeline.py +++ b/sdks/python/apache_beam/pipeline.py @@ -1081,8 +1081,6 @@ def to_runner_api( # because other runners do not have this option. context.enable_best_effort_deterministic_pickling = getattr( self.runner, 'enable_best_effort_deterministic_pickling', False) - context.enable_stable_identifier_pickling = getattr( - self.runner, 'enable_stable_code_identifier_pickling', False) # The RunnerAPI spec requires certain transforms and side-inputs to have KV # inputs (and corresponding outputs). From 7289fb69a7b7e03c9ece29e16047f404d8786247 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 13:01:24 -0400 Subject: [PATCH 07/13] Refactor code object identifier and hashing logic Refactor code object handling and improve normalization for hashing. --- .../internal/code_object_pickler.py | 67 ++++++++++++++----- 1 file changed, 49 insertions(+), 18 deletions(-) diff --git a/sdks/python/apache_beam/internal/code_object_pickler.py b/sdks/python/apache_beam/internal/code_object_pickler.py index 269bccb6b461..a2c634ada22f 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler.py +++ b/sdks/python/apache_beam/internal/code_object_pickler.py @@ -51,6 +51,7 @@ import collections import hashlib +import importlib import inspect import re import sys @@ -213,7 +214,7 @@ def _search_function( first_part = qual_name_parts[0] if (node.__code__ == callable.__code__): if len(qual_name_parts) > 1: - raise ValueError('Qual name parts too long') + return None return '__code__' # If first part is '' then the code object is in a local variable # so we should add __code__ to the path to indicate that we are entering @@ -236,15 +237,12 @@ def _search_code( Returns: The code object identifier, or None if not found. - - Raises: - ValueError: If the qual name parts are too long. """ first_part = qual_name_parts[0] rest = qual_name_parts[1:] if hasattr(callable, '__code__') and node == callable.__code__: if len(qual_name_parts) > 1: - raise ValueError('Qual name parts too long') + return None return '' elif first_part == '': code_objects_by_name = collections.defaultdict(list) @@ -255,12 +253,11 @@ def _search_code( # If there is only one lambda, we can use the default path # 'co_consts[]'. This is the most common case and it is # faster than calculating the signature and the hash. - if num_lambdas == 1: + if num_lambdas == 1 and rest and rest[0] == '': path = _search(callable, code_objects_by_name[''][0], rest) if path is not None: return _extend_path('co_consts[]', path) - else: - return _search_lambda(callable, code_objects_by_name, rest) + return _search_lambda(callable, code_objects_by_name, rest) elif node.co_name == first_part: return _search(callable, node, rest) @@ -281,6 +278,8 @@ def _search_lambda( """ # There are multiple lambdas in the code object, so we need to calculate # the signature and the hash to identify the correct lambda. + if not qual_name_parts: + return None lambda_code_objects_by_name = collections.defaultdict(list) name = qual_name_parts[0] code_objects = code_objects_by_name[name] @@ -306,24 +305,25 @@ def _search_lambda( else: # If there is only one lambda with this signature, we can # use the signature to identify the correct lambda. - path = _search(callable, code_objects[0], qual_name_parts) + path = _search(callable, lambda_objects[0], qual_name_parts) if path is not None: return _extend_path(f'co_consts[{lambda_name}]', path) else: # For non lambda objects, we can use the name to identify the object. - path = _search(callable, code_objects[0], qual_name_parts) - if path is not None: - return _extend_path(f'co_consts[{name}]', path) + for code_object in code_objects: + path = _search(callable, code_objects[0], qual_name_parts) + if path is not None: + return _extend_path(f'co_consts[{name}]', path) # Matches a path like: co_consts[my_function] _SINGLE_NAME_PATTERN = re.compile(r'co_consts\[([a-zA-Z0-9\<\>_-]+)]') -# Matches a path like: co_consts[, ('x',)] +# Matches a path like: co_consts[, ()] or co_consts[, ('x',)] _LAMBDA_WITH_ARGS_PATTERN = re.compile( - r"co_consts\[(<.*?>),\s(\('[^']+'(?:,\s*'[^']+')*,?\))\]") + r'co_consts\[(<[^>]+>),\s*(\([^\)]*\))\]') # Matches a path like: co_consts[, ('x',), 1234567890] _LAMBDA_WITH_HASH_PATTERN = re.compile( - r"co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]") + r'co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]') # Matches a path like: __defaults__[0] _DEFAULT_PATTERN = re.compile(r'(__defaults__)\[(\d+)\]') # Matches an argument like: 'x' @@ -435,10 +435,24 @@ def get_code_from_identifier(code_object_identifier: str): if not code_object_identifier: raise ValueError('Path must not be empty.') parts = code_object_identifier.split('.') - if parts[0] not in sys.modules: + obj = None + start_index = 0 + for i in range(len(parts), 0, -1): + module_name = '.'.join(parts[:i]) + if module_name in sys.modules and sys.modules[module_name] is not None: + obj = sys.modules[module_name] + start_index = i + break + try: + obj = importlib.import_module(module_name) + start_index = i + break + except (ImportError, ModuleNotFoundError, ValueError): + continue + if obj is None: raise AttributeError(f'Module {parts[0]} not found in sys.modules') obj = sys.modules[parts[0]] - for part in parts[1:]: + for part in parts[start_index]: if name_result := _SINGLE_NAME_PATTERN.fullmatch(part): obj = _get_code_object_from_single_name_pattern( obj, name_result, code_object_identifier) @@ -495,6 +509,15 @@ def _get_arg_count(obj: types.CodeType): + (obj.co_flags & 8 == 8) # PyCF_VARKEYWORDS ) +def _normalize_const(const): + """Normalizes a constant from co_consts for deterministic hashing.""" + if inspect.iscode(const): + return ('__code__', _create_bytecode_hash(const)) + if isinstance(const, tuple): + return tuple(_normalize_const(x) for x in const) + if isinstance(const, frozenset): + return frozenset(_normalize_const(x) for x in const) + return (type(const).__qualname__, repr(const)) def _create_bytecode_hash(code_object: types.CodeType): """Returns the hash of a code object. @@ -505,4 +528,12 @@ def _create_bytecode_hash(code_object: types.CodeType): Returns: The hash of the code object. """ - return hashlib.md5(code_object.co_code).hexdigest() + hash_input = repr(( + code_object.co_code, + tuple(_normalize_const(c) for c in code_object.co_consts), + code_object.co_names, + code_object.co_varnames, + code_object.co_freevars, + code_object.co_cellvars, + )).encode('utf-8') + return hashlib.md5(hash_input).hexdigest() From 195bc3522a64484f326d6fbc5298247765a3f1c5 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 13:06:16 -0400 Subject: [PATCH 08/13] Enhance tests for lambda and decorator functionality Added new test cases for lambda functions and decorators in code_object_pickler. --- .../internal/code_object_pickler_test.py | 146 +++++++++++++++++- 1 file changed, 142 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/internal/code_object_pickler_test.py b/sdks/python/apache_beam/internal/code_object_pickler_test.py index abe404ff02c5..b677590cf5e4 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler_test.py +++ b/sdks/python/apache_beam/internal/code_object_pickler_test.py @@ -17,7 +17,9 @@ """Tests for generating stable identifiers to use for Pickle serialization.""" -import hashlib +import functools +import sys +import types import unittest from parameterized import parameterized @@ -63,6 +65,9 @@ def get_lambda_from_dictionary_same_args(): d = {"a": lambda x: 1, "b": lambda x: x + 1} return d["a"] +def get_lambda_from_dictionary_no_args(): + d = {"a": lambda: 1, "b": lambda y: 2} + return d["a"] def function_with_lambda_default_argument(fn=lambda x: 1): return fn @@ -75,6 +80,41 @@ def function_with_function_default_argument(fn=top_level_function): def function_decorator(f): return lambda x: f(f(x)) +def get_zero_arg_and_one_arg_lambdas(): + return (lambda: 1, lambda x: 2) + + +def get_same_co_code_lambdas(): + tag_a = "train" + tag_b = "test" + return (lambda x: (tag_a, x), lambda x: (tag_b, x)) + + +def get_single_lambda_and_nested_function_with_lambda(): + unused_outer_lambda = lambda x: x + 1 + + def inner_fn(): + return lambda y: y + 2 + + return inner_fn() + + +def shared_wrapper_decorator(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + return wrapper + + +@shared_wrapper_decorator +def outer_decorated_with_shared_wrapper(): + @shared_wrapper_decorator + def inner_decorated(): + return 1 + + return inner_decorated + @function_decorator def add_one(x): @@ -143,15 +183,50 @@ def get_lambda_from_dictionary(): f"{prefix}" ".get_lambda_from_dictionary.__code__.co_consts[, ('x',)]") ), + ( + get_lambda_from_dictionary_no_args(), + ( + f"{prefix}" + ".get_lambda_from_dictionary_no_args.__code__.co_consts" + "[, ()]" + ), + ), ( get_lambda_from_dictionary_same_args(), ( f"{prefix}" ".get_lambda_from_dictionary_same_args.__code__.co_consts" - "[, ('x',), " + hashlib.md5( - get_lambda_from_dictionary_same_args().__code__.co_code). - hexdigest() + "]")), + "[, ('x',), " + + code_object_pickler._create_bytecode_hash( + get_lambda_from_dictionary_same_args().__code__ + ) + + "]" + ), + ), + ( + get_zero_arg_and_one_arg_lambdas()[0], + ( + f"{prefix}" + ".get_zero_arg_and_one_arg_lambdas.__code__.co_consts" + "[, ()]" + ), + ), ( + get_zero_arg_and_one_arg_lambdas()[1], + ( + f"{prefix}" + ".get_zero_arg_and_one_arg_lambdas.__code__.co_consts" + "[, ('x',)]" + ), + ), + ( + get_single_lambda_and_nested_function_with_lambda(), + ( + f"{prefix}" + ".get_single_lambda_and_nested_function_with_lambda.__code__" + ".co_consts[inner_fn].co_consts[]" + ), + ), function_with_lambda_default_argument(), ( f"{prefix}" @@ -214,6 +289,69 @@ def test_roundtrip(self, callable, unused_path): actual = code_object_pickler.get_code_from_identifier(path) self.assertEqual(actual, callable.__code__) + def test_two_lambdas_where_one_has_no_arguments_roundtrips(self): + zero_arg_lambda, one_arg_lambda = get_zero_arg_and_one_arg_lambdas() + zero_arg_path = code_object_pickler.get_code_object_identifier( + zero_arg_lambda + ) + one_arg_path = code_object_pickler.get_code_object_identifier( + one_arg_lambda + ) + + self.assertEqual( + zero_arg_path, + f"{prefix}.get_zero_arg_and_one_arg_lambdas.__code__.co_consts" + "[, ()]", + ) + self.assertEqual( + one_arg_path, + f"{prefix}.get_zero_arg_and_one_arg_lambdas.__code__.co_consts" + "[, ('x',)]", + ) + self.assertEqual( + code_object_pickler.get_code_from_identifier(zero_arg_path), + zero_arg_lambda.__code__, + ) + self.assertEqual( + code_object_pickler.get_code_from_identifier(one_arg_path), + one_arg_lambda.__code__, + ) + + def test_same_co_code_lambdas_roundtrip_distinct_code_objects(self): + lambda_a, lambda_b = get_same_co_code_lambdas() + path_a = code_object_pickler.get_code_object_identifier(lambda_a) + path_b = code_object_pickler.get_code_object_identifier(lambda_b) + + self.assertNotEqual(path_a, path_b) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_a), + lambda_a.__code__, + ) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_b), + lambda_b.__code__, + ) + + def test_shared_decorator_wrapper_on_outer_and_inner_returns_none(self): + inner = outer_decorated_with_shared_wrapper() + self.assertIsNone(code_object_pickler.get_code_object_identifier(inner)) + + def test_module_reexporting_same_named_leaf_resolves_longest_module(self): + parent_mod = types.ModuleType("_fake_pkg") + child_mod = types.ModuleType("_fake_pkg.leaf") + child_mod.leaf = top_level_function + # Simulate `from _fake_pkg.leaf import leaf` in `_fake_pkg/__init__.py` + parent_mod.leaf = top_level_function + sys.modules["_fake_pkg"] = parent_mod + sys.modules["_fake_pkg.leaf"] = child_mod + try: + actual = code_object_pickler.get_code_from_identifier( + "_fake_pkg.leaf.leaf.__code__" + ) + self.assertEqual(actual, top_level_function.__code__) + finally: + sys.modules.pop("_fake_pkg", None) + sys.modules.pop("_fake_pkg.leaf", None) class GetCodeFromCodeObjectIdentifierTest(unittest.TestCase): def test_empty_path_raises_exception(self): From 48d8d4e85b7cb4c8c1c5e428de58f58ce8843f47 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 14:17:43 -0400 Subject: [PATCH 09/13] Refactor code object search and lambda handling Refactor code object search logic to handle local code objects and update lambda handling. Replace lambda-specific functions with more general object handling functions. --- .../internal/code_object_pickler.py | 125 +++++++++--------- 1 file changed, 60 insertions(+), 65 deletions(-) diff --git a/sdks/python/apache_beam/internal/code_object_pickler.py b/sdks/python/apache_beam/internal/code_object_pickler.py index a2c634ada22f..09c78ec96012 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler.py +++ b/sdks/python/apache_beam/internal/code_object_pickler.py @@ -249,24 +249,16 @@ def _search_code( for co_const in node.co_consts: if inspect.iscode(co_const): code_objects_by_name[co_const.co_name].append(co_const) - num_lambdas = len(code_objects_by_name.get('', [])) - # If there is only one lambda, we can use the default path - # 'co_consts[]'. This is the most common case and it is - # faster than calculating the signature and the hash. - if num_lambdas == 1 and rest and rest[0] == '': - path = _search(callable, code_objects_by_name[''][0], rest) - if path is not None: - return _extend_path('co_consts[]', path) - return _search_lambda(callable, code_objects_by_name, rest) + return _search_locals(callable, code_objects_by_name, rest) elif node.co_name == first_part: return _search(callable, node, rest) -def _search_lambda( +def _search_locals( callable: types.FunctionType, code_objects_by_name: dict[str, list[types.CodeType]], qual_name_parts: list[str]): - """Searches a lambda to create a code object identifier. + """Searches local code objects to create a code object identifier. Args: callable: The callable to search for. @@ -276,54 +268,57 @@ def _search_lambda( Returns: The code object identifier, or None if not found. """ - # There are multiple lambdas in the code object, so we need to calculate - # the signature and the hash to identify the correct lambda. if not qual_name_parts: return None - lambda_code_objects_by_name = collections.defaultdict(list) name = qual_name_parts[0] code_objects = code_objects_by_name[name] if not code_objects: return None - if name == '': - for code_object in code_objects: - lambda_name = f', {_signature(code_object)}' - lambda_code_objects_by_name[lambda_name].append(code_object) - # Check if there are any lambdas with the same signature. - # If there are, we need to calculate the hash to identify the correct - # lambda. - for lambda_name, lambda_objects in lambda_code_objects_by_name.items(): - if len(lambda_objects) > 1: - for lambda_object in lambda_objects: - path = _search(callable, lambda_object, qual_name_parts) - if path is not None: - return _extend_path( - f'co_consts[{lambda_name},' - f' {_create_bytecode_hash(lambda_object)}]', - path, - ) - else: - # If there is only one lambda with this signature, we can - # use the signature to identify the correct lambda. - path = _search(callable, lambda_objects[0], qual_name_parts) + # If there is only one code object with this name, we can use the default path + # 'co_consts[]'. This is the most common case and it is faster than + # calculating the signature and the hash. + if len(code_objects) == 1: + path = _search(callable, code_objects[0], qual_name_parts) + if path is not None: + return _extend_path(f'co_consts[{name}]', path) + return None + + # There are multiple code objects with the same name, so we need to calculate + # the signature and the hash to identify the correct code object. + code_objects_by_sig = collections.defaultdict(list) + for code_object in code_objects: + sig_name = f'{name}, {_signature(code_object)}' + code_objects_by_sig[sig_name].append(code_object) + # Check if there are any code objects with the same signature. + # If there are, we need to calculate the hash to identify the correct + # code object. + for sig_name, sig_objects in code_objects_by_sig.items(): + if len(sig_objects) > 1: + for code_object in sig_objects: + path = _search(callable, code_object, qual_name_parts) if path is not None: - return _extend_path(f'co_consts[{lambda_name}]', path) - else: - # For non lambda objects, we can use the name to identify the object. - for code_object in code_objects: - path = _search(callable, code_objects[0], qual_name_parts) + return _extend_path( + f'co_consts[{sig_name},' + f' {_create_bytecode_hash(code_object)}]', + path, + ) + else: + # If there is only one code object with this signature, we can + # use the signature to identify the correct code object. + path = _search(callable, sig_objects[0], qual_name_parts) if path is not None: - return _extend_path(f'co_consts[{name}]', path) + return _extend_path(f'co_consts[{sig_name}]', path) # Matches a path like: co_consts[my_function] _SINGLE_NAME_PATTERN = re.compile(r'co_consts\[([a-zA-Z0-9\<\>_-]+)]') -# Matches a path like: co_consts[, ()] or co_consts[, ('x',)] -_LAMBDA_WITH_ARGS_PATTERN = re.compile( - r'co_consts\[(<[^>]+>),\s*(\([^\)]*\))\]') -# Matches a path like: co_consts[, ('x',), 1234567890] -_LAMBDA_WITH_HASH_PATTERN = re.compile( - r'co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]') +# Matches a path like: co_consts[, ()] or co_consts[my_function, ('x',)] +_OBJECT_WITH_ARGS_PATTERN = re.compile( + r'co_consts\[([a-zA-Z0-9\<\>_-]+),\s*(\([^\)]*\))\]') +# Matches a path like: co_consts[, ('x',), 1234567890] or +# co_consts[my_function, ('x',), 1234567890] +_OBJECT_WITH_HASH_PATTERN = re.compile( + r'co_consts\[([a-zA-Z0-9\<\>_-]+),\s*(\([^\)]*\)),\s*(.+)\]') # Matches a path like: __defaults__[0] _DEFAULT_PATTERN = re.compile(r'(__defaults__)\[(\d+)\]') # Matches an argument like: 'x' @@ -357,13 +352,13 @@ def _get_code_object_from_single_name_pattern( raise AttributeError(f'Could not find code object with path: {path}') -def _get_code_object_from_lambda_with_args_pattern( - obj: types.ModuleType, lambda_with_args_result: re.Match[str], path: str): - """Returns the code object from a lambda with args pattern. +def _get_code_object_from_object_with_args_pattern( + obj: types.ModuleType, object_with_args_result: re.Match[str], path: str): + """Returns the code object from an object with args pattern. Args: obj: The object to search within. - lambda_with_args_result: The result of the lambda with args pattern search. + object_with_args_result: The result of the lambda with args pattern search. path: The path to the code object. Returns: @@ -372,7 +367,7 @@ def _get_code_object_from_lambda_with_args_pattern( Raises: AttributeError: If the code object is not found. """ - name = lambda_with_args_result.group(1) + name = object_with_args_result.group(1) code_objects = collections.defaultdict(list) if hasattr(obj, 'co_consts'): for co_const in obj.co_consts: @@ -381,19 +376,19 @@ def _get_code_object_from_lambda_with_args_pattern( for name, objects in code_objects.items(): for obj_ in objects: args = tuple( - re.findall(_ARGUMENT_PATTERN, lambda_with_args_result.group(2))) + re.findall(_ARGUMENT_PATTERN, object_with_args_result.group(2))) if obj_.co_varnames[:_get_arg_count(obj_)] == args: return obj_ raise AttributeError(f'Could not find code object with path: {path}') -def _get_code_object_from_lambda_with_hash_pattern( - obj: types.ModuleType, lambda_with_hash_result: re.Match[str], path: str): +def _get_code_object_from_object_with_hash_pattern( + obj: types.ModuleType, object_with_hash_result: re.Match[str], path: str): """Returns the code object from a lambda with hash pattern. Args: obj: The object to search within. - lambda_with_hash_result: The result of the lambda with hash pattern search. + object_with_hash_result: The result of the lambda with hash pattern search. path: The path to the code object. Returns: @@ -402,7 +397,7 @@ def _get_code_object_from_lambda_with_hash_pattern( Raises: AttributeError: If the code object is not found. """ - name = lambda_with_hash_result.group(1) + name = object_with_hash_result.group(1) code_objects = collections.defaultdict(list) if hasattr(obj, 'co_consts'): for co_const in obj.co_consts: @@ -411,9 +406,9 @@ def _get_code_object_from_lambda_with_hash_pattern( for name, objects in code_objects.items(): for obj_ in objects: args = tuple( - re.findall(_ARGUMENT_PATTERN, lambda_with_hash_result.group(2))) + re.findall(_ARGUMENT_PATTERN, object_with_hash_result.group(2))) if obj_.co_varnames[:_get_arg_count(obj_)] == args: - hash_value = lambda_with_hash_result.group(3) + hash_value = object_with_hash_result.group(3) if hash_value == str(_create_bytecode_hash(obj_)): return obj_ raise AttributeError(f'Could not find code object with path: {path}') @@ -456,12 +451,12 @@ def get_code_from_identifier(code_object_identifier: str): if name_result := _SINGLE_NAME_PATTERN.fullmatch(part): obj = _get_code_object_from_single_name_pattern( obj, name_result, code_object_identifier) - elif lambda_with_args_result := _LAMBDA_WITH_ARGS_PATTERN.fullmatch(part): - obj = _get_code_object_from_lambda_with_args_pattern( - obj, lambda_with_args_result, code_object_identifier) - elif lambda_with_hash_result := _LAMBDA_WITH_HASH_PATTERN.fullmatch(part): - obj = _get_code_object_from_lambda_with_hash_pattern( - obj, lambda_with_hash_result, code_object_identifier) + elif object_with_args_result := _OBJECT_WITH_ARGS_PATTERN.fullmatch(part): + obj = _get_code_object_from_object_with_args_pattern( + obj, object_with_args_result, code_object_identifier) + elif object_with_hash_result := _OBJECT_WITH_HASH_PATTERN.fullmatch(part): + obj = _get_code_object_from_object_with_hash_pattern( + obj, object_with_hash_result, code_object_identifier) elif default_result := _DEFAULT_PATTERN.fullmatch(part): index = int(default_result.group(2)) if index >= len(obj.__defaults__): From 5dfd47bd6bf5b32d53b4a05e736e8c3fbcb09bc7 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 14:33:35 -0400 Subject: [PATCH 10/13] Add tests for same-named nested functions Added tests for handling same-named nested functions with different and same arguments in code object pickler. --- .../internal/code_object_pickler_test.py | 94 +++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/sdks/python/apache_beam/internal/code_object_pickler_test.py b/sdks/python/apache_beam/internal/code_object_pickler_test.py index b677590cf5e4..92077add2800 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler_test.py +++ b/sdks/python/apache_beam/internal/code_object_pickler_test.py @@ -115,6 +115,27 @@ def inner_decorated(): return inner_decorated +def get_same_named_nested_functions_different_args(branch): + if branch: + def nested_fn(x): + return x + 1 + else: + def nested_fn(x, y): + return x + y + + return nested_fn + + +def get_same_named_nested_functions_same_args(branch): + if branch: + def nested_fn(x): + return x + 1 + else: + def nested_fn(x): + return x + 2 + + return nested_fn + @function_decorator def add_one(x): @@ -226,6 +247,46 @@ def get_lambda_from_dictionary(): ".get_single_lambda_and_nested_function_with_lambda.__code__" ".co_consts[inner_fn].co_consts[]" ), + ), + ( + get_same_named_nested_functions_different_args(True), + ( + f"{prefix}" + ".get_same_named_nested_functions_different_args.__code__.co_consts" + "[nested_fn, ('x',)]" + ), + ), + ( + get_same_named_nested_functions_different_args(False), + ( + f"{prefix}" + ".get_same_named_nested_functions_different_args.__code__.co_consts" + "[nested_fn, ('x', 'y')]" + ), + ), + ( + get_same_named_nested_functions_same_args(True), + ( + f"{prefix}" + ".get_same_named_nested_functions_same_args.__code__.co_consts" + "[nested_fn, ('x',), " + + code_object_pickler._create_bytecode_hash( + get_same_named_nested_functions_same_args(True).__code__ + ) + + "]" + ), + ), + ( + get_same_named_nested_functions_same_args(False), + ( + f"{prefix}" + ".get_same_named_nested_functions_same_args.__code__.co_consts" + "[nested_fn, ('x',), " + + code_object_pickler._create_bytecode_hash( + get_same_named_nested_functions_same_args(False).__code__ + ) + + "]" + ), ), function_with_lambda_default_argument(), ( @@ -353,6 +414,39 @@ def test_module_reexporting_same_named_leaf_resolves_longest_module(self): sys.modules.pop("_fake_pkg", None) sys.modules.pop("_fake_pkg.leaf", None) +def test_same_named_nested_functions_with_different_args_roundtrip(self): + fn_true = get_same_named_nested_functions_different_args(True) + fn_false = get_same_named_nested_functions_different_args(False) + path_true = code_object_pickler.get_code_object_identifier(fn_true) + path_false = code_object_pickler.get_code_object_identifier(fn_false) + + self.assertNotEqual(path_true, path_false) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_true), + fn_true.__code__, + ) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_false), + fn_false.__code__, + ) + + def test_same_named_nested_functions_with_same_args_roundtrip(self): + fn_true = get_same_named_nested_functions_same_args(True) + fn_false = get_same_named_nested_functions_same_args(False) + path_true = code_object_pickler.get_code_object_identifier(fn_true) + path_false = code_object_pickler.get_code_object_identifier(fn_false) + + self.assertNotEqual(path_true, path_false) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_true), + fn_true.__code__, + ) + self.assertEqual( + code_object_pickler.get_code_from_identifier(path_false), + fn_false.__code__, + ) + + class GetCodeFromCodeObjectIdentifierTest(unittest.TestCase): def test_empty_path_raises_exception(self): with self.assertRaisesRegex(ValueError, "Path must not be empty"): From 06527197fe990e5ff47673c80dc4dc1d118fffb8 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 15:41:31 -0400 Subject: [PATCH 11/13] Add parentheses --- sdks/python/apache_beam/internal/code_object_pickler_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/python/apache_beam/internal/code_object_pickler_test.py b/sdks/python/apache_beam/internal/code_object_pickler_test.py index 92077add2800..0e5ca41be953 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler_test.py +++ b/sdks/python/apache_beam/internal/code_object_pickler_test.py @@ -288,6 +288,7 @@ def get_lambda_from_dictionary(): + "]" ), ), + ( function_with_lambda_default_argument(), ( f"{prefix}" From 046921981cb785adb15d33f9fe9000e83c006285 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 15:43:17 -0400 Subject: [PATCH 12/13] Fix formatting --- sdks/python/apache_beam/internal/code_object_pickler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/python/apache_beam/internal/code_object_pickler.py b/sdks/python/apache_beam/internal/code_object_pickler.py index 09c78ec96012..9da0489821fd 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler.py +++ b/sdks/python/apache_beam/internal/code_object_pickler.py @@ -504,6 +504,7 @@ def _get_arg_count(obj: types.CodeType): + (obj.co_flags & 8 == 8) # PyCF_VARKEYWORDS ) + def _normalize_const(const): """Normalizes a constant from co_consts for deterministic hashing.""" if inspect.iscode(const): @@ -514,6 +515,7 @@ def _normalize_const(const): return frozenset(_normalize_const(x) for x in const) return (type(const).__qualname__, repr(const)) + def _create_bytecode_hash(code_object: types.CodeType): """Returns the hash of a code object. From 347f28f0e952b281bebee666fdeaf0f950feedc3 Mon Sep 17 00:00:00 2001 From: kristynsmith Date: Wed, 21 Oct 2026 16:00:50 -0400 Subject: [PATCH 13/13] Fix indentation for test function name --- sdks/python/apache_beam/internal/code_object_pickler_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/internal/code_object_pickler_test.py b/sdks/python/apache_beam/internal/code_object_pickler_test.py index 0e5ca41be953..9a520b16e70d 100644 --- a/sdks/python/apache_beam/internal/code_object_pickler_test.py +++ b/sdks/python/apache_beam/internal/code_object_pickler_test.py @@ -431,7 +431,7 @@ def test_same_named_nested_functions_with_different_args_roundtrip(self): fn_false.__code__, ) - def test_same_named_nested_functions_with_same_args_roundtrip(self): +def test_same_named_nested_functions_with_same_args_roundtrip(self): fn_true = get_same_named_nested_functions_same_args(True) fn_false = get_same_named_nested_functions_same_args(False) path_true = code_object_pickler.get_code_object_identifier(fn_true)