kristynsmith commented on code in PR #40711:
URL: https://github.com/apache/beam/pull/40711#discussion_r4473782999
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -279,51 +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.
- lambda_code_objects_by_name = collections.defaultdict(list)
+ if not qual_name_parts:
+ return None
name = qual_name_parts[0]
code_objects = code_objects_by_name[name]
if not code_objects:
return None
- if name == '<lambda>':
- for code_object in code_objects:
- lambda_name = f'<lambda>, {_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, code_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.
+ # If there is only one code object with this name, we can use the default
path
+ # 'co_consts[<name>]'. 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[{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[{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[<lambda>, ('x',)]
-_LAMBDA_WITH_ARGS_PATTERN = re.compile(
- r"co_consts\[(<.*?>),\s(\('[^']+'(?:,\s*'[^']+')*,?\))\]")
-# Matches a path like: co_consts[<lambda>, ('x',), 1234567890]
-_LAMBDA_WITH_HASH_PATTERN = re.compile(
- r"co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]")
+# Matches a path like: co_consts[<lambda>, ()] or co_consts[my_function,
('x',)]
+_OBJECT_WITH_ARGS_PATTERN = re.compile(
Review Comment:
good point, i didn't think about non-ASCII characters
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -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.
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]
