Merge pull request #56 from LSSTDESC/negative-flag

Negative flag
LSSTDESC · Jul 19, 2021 · 1ce4fc8 · 1ce4fc8
2 parents 836a16f + 79f2ed7
commit 1ce4fc8
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 3 deletions.
diff --git a/ceci/stage.py b/ceci/stage.py
@@ -303,8 +303,7 @@ def main(cls):
         return 0
 
     @classmethod
-    def _parse_command_line(cls):
-        cmd = " ".join(sys.argv[:])
+    def _parse_command_line(cls, cmd=None):
         import argparse
 
         parser = argparse.ArgumentParser(description=f"Run pipeline stage {cls.name}")
@@ -315,6 +314,7 @@ def _parse_command_line(cls):
 
             if opt_type == bool:
                 parser.add_argument(f"--{conf}", action="store_const", const=True)
+                parser.add_argument(f"--no-{conf}", dest=conf, action="store_const", const=False)
             elif opt_type == list:
                 out_type = def_val[0] if type(def_val[0]) == type else type(def_val[0])
                 if out_type is str:
@@ -363,7 +363,12 @@ def _parse_command_line(cls):
             default=0,
             help="Report memory use. Argument gives interval in seconds between reports",
         )
-        args = parser.parse_args()
+
+        if cmd is None:
+            args = parser.parse_args()
+        else:
+            args = parser.parse_args(cmd)
+
         return args
 
     @classmethod

diff --git a/tests/test_stage.py b/tests/test_stage.py
@@ -214,6 +214,73 @@ class Golf(PipelineStage):
             def run(self):
                 pass
 
+def test_bool_flags():
+    class Hotel(PipelineStage):
+        inputs = []
+        outputs = []
+        config_options = {'xyz': bool}
+
+    cmd = ["Hotel", "--config", "tests/config.yml"]
+
+    # Basic case with a single flag
+    h = Hotel(Hotel._parse_command_line(cmd + ["--xyz"]))
+    assert h.config['xyz'] is True
+
+    h = Hotel(Hotel._parse_command_line(cmd + ["--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    # check latter takes precedence if both specified
+    h = Hotel(Hotel._parse_command_line(cmd + ["--xyz", "--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    h = Hotel(Hotel._parse_command_line(cmd + ["--no-xyz", "--xyz"]))
+    assert h.config['xyz'] is True
+
+    # flag is not optional here so must be set
+    with pytest.raises(ValueError):
+        h = Hotel(Hotel._parse_command_line(cmd))
+
+    class Hotel2(PipelineStage):
+        inputs = []
+        outputs = []
+        config_options = {'xyz': False}
+
+    h = Hotel2(Hotel2._parse_command_line(cmd + ["--xyz"]))
+    assert h.config['xyz'] is True
+
+    h = Hotel2(Hotel2._parse_command_line(cmd + ["--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    h = Hotel2(Hotel2._parse_command_line(cmd + ["--xyz", "--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    h = Hotel2(Hotel2._parse_command_line(cmd + ["--no-xyz", "--xyz"]))
+    assert h.config['xyz'] is True
+
+    h = Hotel2(Hotel2._parse_command_line(cmd))
+    assert h.config['xyz'] is False
+
+    class Hotel3(PipelineStage):
+        inputs = []
+        outputs = []
+        config_options = {'xyz': True}
+
+    h = Hotel3(Hotel3._parse_command_line(cmd + ["--xyz"]))
+    assert h.config['xyz'] is True
+
+    h = Hotel3(Hotel3._parse_command_line(cmd + ["--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    h = Hotel3(Hotel3._parse_command_line(cmd + ["--xyz", "--no-xyz"]))
+    assert h.config['xyz'] is False
+
+    h = Hotel3(Hotel3._parse_command_line(cmd + ["--no-xyz", "--xyz"]))
+    assert h.config['xyz'] is True
+
+    h = Hotel3(Hotel3._parse_command_line(cmd))
+    assert h.config['xyz'] is True
+
+
 
 def test_unknown_stage():
     with pytest.raises(StageNotFound):