tvalentyn commented on code in PR #25743:
URL: https://github.com/apache/beam/pull/25743#discussion_r1140866962
##########
sdks/python/apache_beam/transforms/core.py:
##########
@@ -1387,6 +1388,58 @@ def partition_for(self, element, num_partitions, *args,
**kwargs):
return self._fn(element, num_partitions, *args, **kwargs)
+def _get_function_body_without_inners(func):
+ source_lines = inspect.getsourcelines(func)[0]
+ source_lines = dropwhile(lambda x: x.startswith("@"), source_lines)
+ def_line = next(source_lines).strip()
+ if def_line.startswith("def ") and def_line.endswith(":"):
+ first_line = next(source_lines)
+ indentation = len(first_line) - len(first_line.lstrip())
+ final_lines = [first_line[indentation:]]
+
+ skip_inner_def = False
+ if first_line[indentation:].startswith("def "):
+ skip_inner_def = True
+ for line in source_lines:
+ line_indentation = len(line) - len(line.lstrip())
+
+ if line[indentation:].startswith("def "):
+ skip_inner_def = True
+ continue
+
+ if skip_inner_def and line_indentation == indentation:
+ skip_inner_def = False
+
+ if skip_inner_def and line_indentation > indentation:
+ continue
+ final_lines.append(line[indentation:])
+
+ return "".join(final_lines)
+ else:
+ return def_line.rsplit(":")[-1].strip()
+
+
+def _check_fn_use_yield_and_return(fn):
+ if isinstance(fn, types.BuiltinFunctionType):
+ return False
+ try:
+ source_code = _get_function_body_without_inners(fn)
+ has_yield = False
+ has_return = False
+ for line in source_code.split("\n"):
+ if line.lstrip().startswith("yield ") or line.lstrip().startswith(
+ "yield("):
+ has_yield = True
+ if line.lstrip().startswith("return ") or line.lstrip().startswith(
+ "return("):
+ has_return = True
+ if has_yield and has_return:
+ return True
+ return False
+ except TypeError:
Review Comment:
I suggest broadening the except to include any exception just in case so as
to not fail the job unnecessarily.
--
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]
