kristynsmith commented on code in PR #40711:
URL: https://github.com/apache/beam/pull/40711#discussion_r4473801026
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -435,19 +430,37 @@ 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:
+ max_module_parts = len(parts)
+ for idx, part in enumerate(parts):
+ if part == '__code__' or part.startswith(('co_consts[', '__defaults__[')):
+ max_module_parts = idx
+ break
+ obj = None
+ start_index = 0
+ for i in range(len(parts), 0, -1):
Review Comment:
i understand, fixed
##########
sdks/python/apache_beam/internal/code_object_pickler_test.py:
##########
@@ -214,6 +367,163 @@ 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"
+ "[<lambda>, ()]",
+ )
+ self.assertEqual(
+ one_arg_path,
+ f"{prefix}.get_zero_arg_and_one_arg_lambdas.__code__.co_consts"
+ "[<lambda>, ('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)
+
+ 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__,
+ )
+
+ def test_same_nested_functions_with_complex_consts_roundtrip(self):
+ fn_true = get_same_named_nested_functions_with_complex_consts(True)
+ fn_false = get_same_named_nested_functions_with_complex_consts(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_get_code_from_identifier_imports_unloaded_module(self):
+ mod_name = "apache_beam.internal.test_data.module_1"
+ saved = sys.modules.pop(mod_name, None)
+ try:
+ code = code_object_pickler.get_code_from_identifier(
+ f"{mod_name}.my_function.__code__")
+ self.assertEqual(code, module_1.my_function.__code__)
+ finally:
+ if saved is not None:
+ sys.modules[mod_name] = saved
+ else:
+ sys.modules.pop(mod_name, None)
+
+ def test_closures_with_different_param_kinds_have_distinct_identifiers(self):
+ lambda_pos, lambda_vararg = get_same_named_lambdas_different_param_kinds()
+ path_pos = code_object_pickler.get_code_object_identifier(lambda_pos)
+ path_vararg = code_object_pickler.get_code_object_identifier(lambda_vararg)
+
+ self.asserNotEqual(path_pos, path_vararg)
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
