move custom argparse action to util

legate/tester/args.py

@@ -17,20 +17,12 @@
 """
 from __future__ import annotations
 
-from argparse import Action, ArgumentParser, Namespace
-from typing import (
-    Any,
-    Generic,
-    Iterable,
-    Iterator,
-    Literal,
-    Sequence,
-    TypeVar,
-    Union,
-)
+from argparse import ArgumentParser
+from typing import Literal, Union
 
 from typing_extensions import TypeAlias
 
+from ..util.args import ExtendAction, MultipleChoices
 from . import (
     DEFAULT_CPUS_PER_NODE,
     DEFAULT_GPU_DELAY,
@@ -41,8 +33,6 @@
     FEATURES,
 )
 
-T = TypeVar("T")
-
 PinOptionsType: TypeAlias = Union[
     Literal["partial"],
     Literal["none"],
@@ -56,57 +46,6 @@
 )
 
 
-class MultipleChoices(Generic[T]):
-    """A container that reports True for any item or subset inclusion.
-
-    Parameters
-    ----------
-    choices: Iterable[T]
-        The values to populate the containter.
-
-    Examples
-    --------
-
-    >>> choices = MultipleChoices(["a", "b", "c"])
-
-    >>> "a" in choices
-    True
-
-    >>> ("b", "c") in choices
-    True
-
-    """
-
-    def __init__(self, choices: Iterable[T]) -> None:
-        self.choices = set(choices)
-
-    def __contains__(self, x: Union[T, Iterable[T]]) -> bool:
-        if isinstance(x, (list, tuple)):
-            return set(x).issubset(self.choices)
-        return x in self.choices
-
-    def __iter__(self) -> Iterator[T]:
-        return self.choices.__iter__()
-
-
-class ExtendAction(Action):
-    """A custom argparse action to collect multiple values into a list."""
-
-    def __call__(
-        self,
-        parser: ArgumentParser,
-        namespace: Namespace,
-        values: Union[str, Sequence[Any], None],
-        option_string: Union[str, None] = None,
-    ) -> None:
-        items = getattr(namespace, self.dest, None) or []
-        if isinstance(values, list):
-            items.extend(values)
-        else:
-            items.append(values)
-        setattr(namespace, self.dest, items)
-
-
 #: The argument parser for test.py
 parser = ArgumentParser(
     description="Run the Cunumeric test suite",
@@ -122,9 +61,7 @@ def __call__(
     dest="features",
     action=ExtendAction,
     choices=MultipleChoices(sorted(FEATURES)),
-    # argpase evidently only expects string returns from the type converter
-    # here, but returning a list of strings seems to work in practice
-    type=lambda s: s.split(","),  # type: ignore[return-value, arg-type]
+    type=lambda s: s.split(","),  # type: ignore
     help="Test Legate with features (also via USE_*)",
 )
 

legate/util/args.py

@@ -16,9 +16,19 @@
 
 import sys
 import warnings
-from argparse import ArgumentParser, Namespace
+from argparse import Action, ArgumentParser, Namespace
 from dataclasses import dataclass, fields
-from typing import Any, Iterable, Literal, Sequence, Type, TypeVar, Union
+from typing import (
+    Any,
+    Generic,
+    Iterable,
+    Iterator,
+    Literal,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+)
 
 from typing_extensions import TypeAlias
 
@@ -29,8 +39,10 @@ class _UnsetType:
 
 Unset = _UnsetType()
 
-_T = TypeVar("_T")
-NotRequired = Union[_UnsetType, _T]
+
+T = TypeVar("T")
+
+NotRequired = Union[_UnsetType, T]
 
 
 # https://docs.python.org/3/library/argparse.html#action
@@ -76,6 +88,58 @@ def entries(obj: Any) -> Iterable[tuple[str, Any]]:
             yield (f.name, value)
 
 
+class MultipleChoices(Generic[T]):
+    """A container that reports True for any item or subset inclusion.
+
+    Parameters
+    ----------
+    choices: Iterable[T]
+        The values to populate the containter.
+
+    Examples
+    --------
+
+    >>> choices = MultipleChoices(["a", "b", "c"])
+
+    >>> "a" in choices
+    True
+
+    >>> ("b", "c") in choices
+    True
+
+    """
+
+    def __init__(self, choices: Iterable[T]) -> None:
+        self._choices = set(choices)
+
+    def __contains__(self, x: Union[T, Sequence[T]]) -> bool:
+        if isinstance(x, (list, tuple)):
+            return set(x).issubset(self._choices)
+        return x in self._choices
+
+    def __iter__(self) -> Iterator[T]:
+        return self._choices.__iter__()
+
+
+class ExtendAction(Action, Generic[T]):
+    """A custom argparse action to collect multiple values into a list."""
+
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Union[str, Sequence[T], None],
+        option_string: Union[str, None] = None,
+    ) -> None:
+        items = getattr(namespace, self.dest) or []
+        if isinstance(values, (list, tuple)):
+            items.extend(values)
+        else:
+            items.append(values)
+        # removing any duplicates before storing
+        setattr(namespace, self.dest, list(set(items)))
+
+
 def parse_library_command_args(
     libname: str, args: Iterable[Argument]
 ) -> Namespace:

tests/unit/legate/tester/stages/test_test_stage.py

@@ -24,7 +24,7 @@
 from legate.tester.config import Config
 from legate.tester.stages import test_stage as m
 from legate.tester.stages.util import StageResult, StageSpec
-from legate.tester.test_system import ProcessResult, TestSystem
+from legate.tester.test_system import ProcessResult, TestSystem as _TestSystem
 
 from . import FakeSystem
 
@@ -39,10 +39,10 @@ class MockTestStage(m.TestStage):
 
     args = ["-foo", "-bar"]
 
-    def __init__(self, config: Config, system: TestSystem) -> None:
+    def __init__(self, config: Config, system: _TestSystem) -> None:
         self._init(config, system)
 
-    def compute_spec(self, config: Config, system: TestSystem) -> StageSpec:
+    def compute_spec(self, config: Config, system: _TestSystem) -> StageSpec:
         return StageSpec(2, [(0,), (1,), (2,)])
 
 

tests/unit/legate/tester/test_args.py

@@ -17,11 +17,6 @@
 """
 from __future__ import annotations
 
-from itertools import chain, combinations
-from typing import Iterable, TypeVar
-
-import pytest
-
 from legate.tester import (
     DEFAULT_CPUS_PER_NODE,
     DEFAULT_GPU_DELAY,
@@ -32,14 +27,6 @@
     args as m,
 )
 
-T = TypeVar("T")
-
-
-# https://docs.python.org/3/library/itertools.html#itertools-recipes
-def powerset(iterable: Iterable[T]) -> Iterable[Iterable[T]]:
-    xs = list(iterable)
-    return chain.from_iterable(combinations(xs, n) for n in range(len(xs) + 1))
-
 
 class TestParserDefaults:
     def test_featurs(self) -> None:
@@ -100,33 +87,3 @@ def test_parser_epilog(self) -> None:
 
     def test_parser_description(self) -> None:
         assert m.parser.description == "Run the Cunumeric test suite"
-
-
-class TestMultipleChoices:
-    @pytest.mark.parametrize("choices", ([1, 2, 3], range(4), ("a", "b")))
-    def test_init(self, choices: Iterable[T]) -> None:
-        mc = m.MultipleChoices(choices)
-        assert mc.choices == set(choices)
-
-    def test_contains_item(self) -> None:
-        choices = [1, 2, 3]
-        mc = m.MultipleChoices(choices)
-        for item in choices:
-            assert item in mc
-
-    def test_contains_subset(self) -> None:
-        choices = [1, 2, 3]
-        mc = m.MultipleChoices(choices)
-        for subset in powerset(choices):
-            assert subset in mc
-
-    def test_iter(self) -> None:
-        choices = [1, 2, 3]
-        mc = m.MultipleChoices(choices)
-        assert list(mc) == choices
-
-
-# Testing this directly would require getting into argparse
-# internals. See test_config.py for indirect tests with --use
-class TestExtendAction:
-    pass

tests/unit/legate/util/test_args.py

@@ -14,13 +14,60 @@
 #
 
 import sys
+from argparse import ArgumentParser
 from dataclasses import dataclass
+from typing import Iterable, TypeVar
 
 import pytest
 
 import legate.util.args as m
 
-from ...util import Capsys
+from ...util import Capsys, powerset
+
+T = TypeVar("T")
+
+
+class TestMultipleChoices:
+    @pytest.mark.parametrize("choices", ([1, 2, 3], range(4), ("a", "b")))
+    def test_init(self, choices: Iterable[T]) -> None:
+        mc = m.MultipleChoices(choices)
+        assert mc._choices == set(choices)
+
+    def test_contains_item(self) -> None:
+        choices = [1, 2, 3]
+        mc = m.MultipleChoices(choices)
+        for item in choices:
+            assert item in mc
+
+    def test_contains_subset(self) -> None:
+        choices = [1, 2, 3]
+        mc = m.MultipleChoices(choices)
+        for subset in powerset(choices):
+            assert subset in mc
+
+    def test_iter(self) -> None:
+        choices = [1, 2, 3]
+        mc = m.MultipleChoices(choices)
+        assert list(mc) == choices
+
+
+class TestExtendAction:
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--foo", dest="foo", action=m.ExtendAction, choices=("a", "b", "c")
+    )
+
+    def test_single(self) -> None:
+        ns = self.parser.parse_args(["--foo", "a"])
+        assert ns.foo == ["a"]
+
+    def test_multi(self) -> None:
+        ns = self.parser.parse_args(["--foo", "a", "--foo", "b"])
+        assert sorted(ns.foo) == ["a", "b"]
+
+    def test_repeat(self) -> None:
+        ns = self.parser.parse_args(["--foo", "a", "--foo", "a"])
+        assert ns.foo == ["a"]
 
 
 @dataclass(frozen=True)