From fddb45b8794f71d568ab8b745d835c4207b2859d Mon Sep 17 00:00:00 2001 From: Yashvardhan Thanvi Date: Sun, 18 Oct 2026 22:58:25 +0530 Subject: [PATCH] [Python] Only give an unknown option a value that directly follows it With retain_unknown_options, get_all_options registers each unknown flag with argparse and parses the flags again. Whether an unknown flag takes a value was decided from the list of unknown args alone, but those are not necessarily adjacent in the original flags. For --unknown_flag --job_name=foo bar the unknown args are ['--unknown_flag', 'bar'], so --unknown_flag was registered as taking a value, and parsing the original flags again failed with "expected one argument". This also made Pipeline.run() fail. Look up where each unknown arg came from in the original flags and only treat an unknown flag as taking a value when that value directly follows it there. A flag written as --flag=value always uses its inline value, so it is no longer registered under the name "--flag=value" when a stray argument follows it. Fixes #22434 --- CHANGES.md | 1 + .../apache_beam/options/pipeline_options.py | 35 +++++++- .../options/pipeline_options_test.py | 88 +++++++++++++++++++ 3 files changed, 123 insertions(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index 04ac961fcd93..5f387e82dafd 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -142,6 +142,7 @@ * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). * (Go) Fixed generic registration (for example `register.Function1x1` or `register.DoFn1x1`) panicking on functions and methods that take or return an anonymous struct. Unnamed struct types are now only registered with the schema registry, since they have no type key for the runtime type registry ([#21789](https://github.com/apache/beam/issues/21789)). +* (Python) An unknown pipeline option is no longer given the value of a later argument that does not directly follow it, which made `Pipeline.run()` and `get_all_options(retain_unknown_options=True)` fail with `expected one argument` for flags like `--unknown_flag --job_name=foo bar` ([#22434](https://github.com/apache/beam/issues/22434)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 5b73b22476b4..4063f7d1458a 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -24,6 +24,7 @@ import json import logging import os +import sys from typing import Any from typing import Callable from typing import Optional @@ -91,6 +92,26 @@ def _f(value): return _f +def _find_positions(args, flags): + """Returns the index in flags of each of args. + + args is expected to be a subsequence of flags, such as the unknown args that + argparse returns from parse_known_args. The position of an arg that cannot be + found is None. + """ + positions = [] + start = 0 + for arg in args: + try: + position = flags.index(arg, start) + except ValueError: + positions.append(None) + continue + positions.append(position) + start = position + 1 + return positions + + class _BeamArgumentParser(argparse.ArgumentParser): """An ArgumentParser that supports ValueProvider options. @@ -535,6 +556,13 @@ def add_new_arg(arg, **kwargs): parser.add_argument(arg, **kwargs) seen.add(arg) + # Unknown args are returned in their original order, but they are not + # necessarily adjacent in the original flags, e.g. a known option may + # sit between an unknown flag and a stray value. Find where each one + # came from so that a value is only given to the flag right before it. + all_flags = sys.argv[1:] if self._flags is None else list(self._flags) + positions = _find_positions(unknown_args, all_flags) + i = 0 while i < len(unknown_args): # End of argument parsing. @@ -545,7 +573,12 @@ def add_new_arg(arg, **kwargs): if not unknown_args[i].startswith('-'): i += 1 continue - if i + 1 >= len(unknown_args) or unknown_args[i + 1].startswith('-'): + has_value = ( + '=' not in unknown_args[i] and i + 1 < len(unknown_args) and + not unknown_args[i + 1].startswith('-') and + positions[i] is not None and positions[i] + 1 < len(all_flags) and + not all_flags[positions[i] + 1].startswith('-')) + if not has_value: split = unknown_args[i].split('=', 1) if len(split) == 1: add_new_arg(unknown_args[i], action='store_true') diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py index 94083a55c6f4..760f7b896add 100644 --- a/sdks/python/apache_beam/options/pipeline_options_test.py +++ b/sdks/python/apache_beam/options/pipeline_options_test.py @@ -407,6 +407,94 @@ def test_retain_unknown_options_unary_single_dash_store_true(self): result = options.get_all_options(retain_unknown_options=True) self.assertEqual(result['i'], True) + # Unknown args are not necessarily adjacent in the original flags, so a + # value must only be attributed to an unknown flag it directly follows. + # See https://github.com/apache/beam/issues/22434. + @parameterized.expand([ + (['--unknown_option', '--job_name=foo', 'bar'], ), + (['--unknown_option', '--job_name', 'foo', 'bar'], ), + ]) + def test_retain_unknown_options_unary_followed_by_known_option(self, flags): + options = PipelineOptions(flags) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertEqual(result['job_name'], 'foo') + + def test_retain_unknown_options_mixed_with_known_options(self): + options = PipelineOptions([ + '--unknown_unary', + '--job_name=foo', + 'stray_value', + '--unknown_binary', + 'some_value', + '--project', + 'bar', + ]) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_unary'], True) + self.assertEqual(result['unknown_binary'], 'some_value') + self.assertEqual(result['job_name'], 'foo') + self.assertEqual(result['project'], 'bar') + + def test_retain_unknown_options_binary_equals_followed_by_value(self): + options = PipelineOptions(['--unknown_option=some_value', 'stray_value']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'some_value') + self.assertNotIn('unknown_option=some_value', result) + + def test_retain_unknown_options_repeated_unary_store_true(self): + options = PipelineOptions( + ['--unknown_option', '--job_name=foo', 'bar', '--unknown_option']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + + def test_retain_unknown_options_repeated_binary_keeps_last_value(self): + options = PipelineOptions([ + '--unknown_option', + 'first', + '--job_name=foo', + '--unknown_option', + 'second', + ]) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'second') + + def test_retain_unknown_options_unary_before_separator(self): + options = PipelineOptions( + ['--unknown_option', '--', 'bar', '--after_separator']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertNotIn('after_separator', result) + + def test_retain_unknown_options_unary_single_dash_followed_by_known_option( + self): + options = PipelineOptions(['-i', '--job_name=foo', 'bar']) + result = options.get_all_options(retain_unknown_options=True) + self.assertEqual(result['i'], True) + self.assertEqual(result['job_name'], 'foo') + + def test_retain_unknown_options_does_not_take_value_of_known_positional(self): + options = PipelineOptions(['--unknown_option', 'some_value']) + + def add_extra_options(parser): + parser.add_argument('extra_positional', nargs='?') + + result = options.get_all_options( + add_extra_args_fn=add_extra_options, retain_unknown_options=True) + self.assertEqual(result['unknown_option'], True) + self.assertEqual(result['extra_positional'], 'some_value') + + def test_retain_unknown_options_value_then_known_positional(self): + options = PipelineOptions(['--unknown_option', 'some_value', 'other']) + + def add_extra_options(parser): + parser.add_argument('extra_positional', nargs='?') + + result = options.get_all_options( + add_extra_args_fn=add_extra_options, retain_unknown_options=True) + self.assertEqual(result['unknown_option'], 'some_value') + self.assertEqual(result['extra_positional'], 'other') + def test_override_options(self): base_flags = ['--num_workers', '5'] options = PipelineOptions(base_flags)