Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Python][Pipeline Options] Recursively Get All Subclasses #31141

Closed
wants to merge 1 commit into from
Closed

[Python][Pipeline Options] Recursively Get All Subclasses #31141

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion sdks/python/apache_beam/options/pipeline_options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,11 @@ def _f(value):

return _f

def get_all_subclasses(cls):
"""Returns all subclasses of a class, recursively."""
return set(cls.__subclasses__()).union(
[s for c in cls.__subclasses__() for s in get_all_subclasses(c)])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this is causing a bunch of test failures because its not correctly deduping items and it considers the current class multiple times. I'm guessing the set comparison isn't working correctly. Could you please take a look?



class _BeamArgumentParser(argparse.ArgumentParser):
"""An ArgumentParser that supports ValueProvider options.
Expand Down Expand Up @@ -326,7 +331,7 @@ def get_all_options(
# instance of each subclass to avoid conflicts.
subset = {}
parser = _BeamArgumentParser()
for cls in PipelineOptions.__subclasses__():
for cls in get_all_subclasses(PipelineOptions):
subset[str(cls)] = cls
for cls in subset.values():
cls._add_argparse_args(parser) # pylint: disable=protected-access
Expand Down
16 changes: 16 additions & 0 deletions sdks/python/apache_beam/options/pipeline_options_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,11 @@ def _add_argparse_args(cls, parser):
parser.add_argument('--option with space', help='mock option with space')
parser.add_argument('--mock_json_option', type=json.loads, default={})

class GrandChildMockOptions(MockOptions):
@classmethod
def _add_argparse_args(cls, parser):
parser.add_argument('--grand_mock_flag', action='store_true', help='mock flag')

# Use with MockOptions in test cases where multiple option classes are needed.
class FakeOptions(PipelineOptions):
@classmethod
Expand Down Expand Up @@ -215,6 +220,17 @@ def test_get_all_options_subclass(self, flags, expected, _):
self.assertEqual(
options.view_as(PipelineOptionsTest.MockOptions).mock_multi_option,
expected['mock_multi_option'])

@parameterized.expand(TEST_CASES)
def test_get_all_options_subsubclass(self, flags, expected, _):
options = PipelineOptionsTest.GrandChildMockOptions()(flags=flags)
self.assertDictContainsSubset(expected, options.get_all_options())
self.assertEqual(
options.view_as(PipelineOptionsTest.GrandChildMockOptions).mock_flag,
expected['mock_flag'])
self.assertEqual(
options.view_as(PipelineOptionsTest.GrandChildMockOptions).grand_mock_flag,
expected['grand_mock_flag'])

@parameterized.expand(TEST_CASES)
def test_get_all_options(self, flags, expected, _):
Expand Down
Loading