Updated chi_parser and chi_run to use sim_state files.

chi_pet/chi_parse.py

@@ -10,6 +10,7 @@
 """
 
 from pathlib import Path
+import warnings
 from . import chi_lib as clib
 
 
@@ -50,13 +51,17 @@ def parse_chi_options():
     create_parser.add_argument('-r', '--replace', default=False, action='store_true',
                                help='Replace simulation file instead of throwing an error if file already exists.(Used with create parser only)')
 
-    create_parser.add_argument('-ny', '--non_yaml', nargs='+', type=Path, default=[],
+    create_parser.add_argument('-ny', '--non-yaml', nargs='+', type=Path, default=[],
                                help='Will add non-yaml files to seed directories when creating directory structure. (Used with create parser only)')
 
     # RUN options
     run_parser = subparsers.add_parser(
         'run', parents=[parent_parser], help='Run a simulation pipeline defined in args yaml file in a singular seed directory. Requires the --args_file option defined.')
 
+    run_parser.add_argument('-us', '--use-sim-states', action='store_true',
+                            default=False,
+                            help='Use sim.state files to determine which states to run.')
+
     # LAUNCH options
     launch_parser = subparsers.add_parser(
         'launch', parents=[parent_parser], help='Launch or create launching script to run simulations in seed directories.')
@@ -71,10 +76,6 @@ def parse_chi_options():
 
     opts = parser.parse_args()
 
-    if opts.command == 'run':
-        if opts.args_file is None:
-            parser.error("'run' requires the '--args_file' option.")
-
     if not opts.workdir:
         opts.workdir = Path.cwd()
     elif not opts.workdir.exists():
@@ -87,7 +88,21 @@ def parse_chi_options():
         with opts.args_file.open('r') as f:
             opts.args_dict = clib.load_yaml_in_order(f)
 
+    if opts.command == 'create':
         if not opts.states:
             opts.states = list(opts.args_dict.keys())
 
+    elif opts.command == 'run':
+        if opts.args_file is None:
+            parser.error("'run' requires the '--args_file' option.")
+
+        if opts.use_sim_states and opts.states:
+            warnings.warn(
+                "Both use_sim_states and states are set. Will prioritize states given in commandline.")
+        elif opts.use_sim_states:  # Only use_sim_states is set
+            for sim_state in opts.workdir.glob('sim.*'):
+                opts.states.append(sim_state.name.split('.')[1])
+        else:  # Neither use_sim_states nor states are set
+            opts.states = list(opts.args_dict.keys())
+
     return opts

chi_pet/chi_run.py

@@ -15,26 +15,28 @@
 class ChiRun(object):
     def __init__(self, opts):
         self._opts = opts
+        self._arg_dict = opts.args_dict
+        self._states = opts.states
 
     def run(self):
         """Run simulation pipeline defined in the the args.yaml file defined as self._opts.args_file.
 
         """
 
         # TODO Add option to see if sim_action file exists and only run those states.
-        for run_state, vals in self._opts.args_dict.items():
+        for run_state, vals in self._args_dict.items():
             args = [str(a) for a in vals]
-            sim_action = Path('sim.{}'.format(run_state))
+            sim_state = Path('sim.{}'.format(run_state))
 
             # If sim action file exists then run the state
-            if run_state in opts.states:
-                if self.run_args(opts.workdir, run_state, args):
-                    (opts.workdir / '.error').touch()
-                elif sim_action.exists():
-                    sim_action.unlink()
-
-    def get_actions(self):
-        return self._opts.args_dict.items()
+            if run_state in self._states:
+                if self.run_args(self._workdir, run_state, args):
+                    (self._opts.workdir / '.error').touch()
+                elif sim_state.exists():
+                    sim_state.unlink()
+
+    def get_run_states(self):
+        return self._states
 
     @classmethod
     def run_args(workdir, state, args):

tests/mhelpers.py

@@ -23,6 +23,8 @@
 MOCK_ARGS_FILE_DICT = {'stage1': ['arg1', 'arg2', 'arg3'],
                        'stage2': ['arg4', 'arg5', 'arg6']}
 
+MOCK_SHELL_ARGS_FILE_DICT = {'touch': ['touch', 'mock_command.txt']}
+
 MOCK_PARAM_DICT_PATH = 'mock_param.yaml'
 MOCK_CHI_PARAM_DICT_PATH = 'mock_param_chi_param.yaml'
 
@@ -51,23 +53,18 @@ def setup_and_teardown(clean_up):
         clean_mocks()  # Clean up after running tests
 
 
-# @pytest.fixture()
-# def mock_non_yaml_file():
-#     ny_file_path = Path.cwd() / 'mock_ny_file.txt'
-#     with ny_file_path.open('w') as nyf:
-#         nyf.write(MOCK_NON_YAML_FILE_STR)
-#     yield ny_file_path
-#     ny_file_path.unlink()
+def mock_non_yaml_file(dir_path):
+    ny_file_path = dir_path / 'mock_ny_file.txt'
+    with ny_file_path.open('w') as nyf:
+        nyf.write(MOCK_NON_YAML_FILE_STR)
+    return ny_file_path
 
 
-@pytest.fixture()
-def mock_args_file(mock_root_dir, clean_up):
-    args_file_path = mock_root_dir / 'mock_args.yaml'
+def mock_args_file(dir_path):
+    args_file_path = dir_path / 'mock_args.yaml'
     with args_file_path.open('w') as aff:
         yaml.dump(MOCK_ARGS_FILE_DICT, aff)
-    yield args_file_path
-    if clean_up:
-        args_file_path.unlink()
+    return args_file_path
 
 
 @pytest.fixture()
@@ -77,6 +74,30 @@ def mock_param_yaml_dict():
     return yaml_dict
 
 
+@pytest.fixture()
+def mock_leaf_dir():
+    """Create a directory to use a leaf chi directory for testing 
+
+    Returns
+    -------
+    Path
+        Path to root directory
+    """
+    chi_leaf_path = Path.cwd() / 'tests/mock_leaf'
+    chi_leaf_path.mkdir()
+    yaml_param_path = chi_leaf_path / MOCK_PARAM_DICT_PATH
+
+    with yaml_param_path.open('w') as ypp:
+        yaml.dump(MOCK_PARAM_DICT, ypp)
+
+    args_file_path = chi_leaf_path / 'mock_args.yaml'
+    with args_file_path.open('w') as aff:
+        yaml.dump(MOCK_ARGS_FILE_DICT, aff)
+    yield args_file_path
+
+    yield chi_leaf_path
+
+
 @pytest.fixture()
 def mock_root_dir():
     """Create a directory to use a root chi directory for testing 
@@ -112,6 +133,9 @@ def opts(x): return None
 def mock_run_opts():
     def opts(x): return None
     opts.command = 'run'
+    opts.use_sim_states = 'False'
     opts.replace = 'False'
     opts.non_yaml = []
+    opts.args_dict = MOCK_ARGS_FILE_DICT
+    opts.states = list(opts.args_dict.keys())
     yield opts

tests/test_chi_node.py

@@ -43,15 +43,14 @@ def test_chi_node_dir_creation(mock_chi_node):
     assert Path('tests/mock_node/data').exists()
 
 
-def test_chi_node_subnode_creation(mock_root_dir, mock_args_file, mock_create_opts):
+def test_chi_node_subnode_creation(mock_root_dir, mock_create_opts):
     """!Test to make sure subnodes are generated properly for basic usage.
     """
     root_path = mock_root_dir
-    mock_args_path = mock_args_file
+    mock_args_path = mock_args_file(root_path)
 
     ny_file_path = root_path / 'mock_ny_file.txt'
-    with ny_file_path.open('w') as nyf:
-        nyf.write(MOCK_NON_YAML_FILE_STR)
+    mock_non_yaml_file(root_path)
 
     mock_create_opts.args_files = [mock_args_path]
 

tests/test_chi_parse.py

@@ -16,8 +16,9 @@
 # ])
 
 
-def test_chi_create_argument_parsing(mock_args_file):
-    mock_args_path = mock_args_file
+def test_chi_create_argument_parsing(mock_root_dir):
+    root_path = mock_root_dir
+    mock_args_path = mock_args_file(root_path)
     # Setup
     sys.argv = ['chi', 'create', 'param1.yaml',
                 'param2.yaml', '-a', f'{mock_args_path}']
@@ -29,8 +30,9 @@ def test_chi_create_argument_parsing(mock_args_file):
         assert Path(pf) in opts.param_files
 
 
-def test_chi_run_argument_parsing(mock_args_file):
-    mock_args_path = mock_args_file
+def test_chi_run_argument_parsing(mock_root_dir):
+    root_path = mock_root_dir
+    mock_args_path = mock_args_file(root_path)
 
     # Setup
     sys.argv = ['chi', 'run', '-a', f'{mock_args_path}']
@@ -40,3 +42,7 @@ def test_chi_run_argument_parsing(mock_args_file):
     sys.argv = ['chi', 'run']
     with pytest.raises(SystemExit):
         opts = parse_chi_options()
+
+# TODO Add test for non-yaml files
+
+# TODO Add test for sim states

tests/test_chi_run.py

@@ -7,13 +7,22 @@
 import pytest
 
 
-# TODO NEXT TEST add tests for chi_run.py
-def test_chi_run_pipeline(mock_run_opts):
+def test_chi_run_states(mock_run_opts):
+    """Test chi_run."""
+    # Make a ChiRun object
+    crun = ChiRun(mock_run_opts)
+
+    assert crun.get_run_states() == mock_run_opts.states
+
+
+def test_chi_run_with_touch(mock_run_opts):
     """Test chi_run."""
-    mock_run_opts.args_dict = MOCK_ARGS_FILE_DICT
     # Make a ChiRun object
     crun = ChiRun(mock_run_opts)
 
-    assert True
+    assert crun.get_run_states() == mock_run_opts.states
+
     # Run the pipeline
     # Make sure action is carried out
+
+# TODO NEXT TEST add tests for chi_run.py