From c428efc8b9a7eb5669b631fe093802e8b5454825 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 44e7568cb8ee..2624185256e6 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -80,6 +80,7 @@ ## Bugfixes +* (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 113dcb00e190..7445551029dd 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 79ea0aa650f8..387e81d4426e 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)