[Python] Only give an unknown option a value that directly follows it

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

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')

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)