ECS-6433 heuristic loading #231

docs/source/upcoming_release_notes/348-ECS-6433_heuristic_loading.rst

@@ -0,0 +1,24 @@
+348 ECS-6433 heuristic loading
+#################
+
+API Breaks
+----------
+- N/A
+
+Features
+--------
+- Updated happi's ``load()`` function to support searching and loading in one user-friendly expression. The load function gathers the results of each search term individually and loads the union of the results.
+
+Bugfixes
+--------
+- N/A
+
+Maintenance
+-----------
+- Cleaned up code by removing stray ``print()`` statements.
+- Added unit tests for the new feature.
+- In the load function of cli.py removed ``term=[]`` and replaced with ``[item]`` for readability.
+
+Contributors
+------------
+- janeliu-slac

happi/cli.py

@@ -396,21 +396,43 @@ def edit(ctx, name: str, edits: list[str]):
 
 
 @happi_cli.command()
-@click.argument('item_names', nargs=-1)
+@click.option('--glob/--regex', 'use_glob', default=True,
+              help='use glob style (default) or regex style search terms. '
+              r'Regex requires backslashes to be escaped (eg. at\\d.\\d)')
+@click.argument('search_criteria', nargs=-1)
 @click.pass_context
-def load(ctx, item_names: list[str]):
-    """Open IPython terminal with ITEM_NAMES loaded."""
+def load(
+    ctx: click.Context,
+    use_glob: bool,
+    search_criteria: list[str]
+):
+    results = set()
+
+    for item in search_criteria:
+        final_results = search_parser(
+            client=get_happi_client_from_config(ctx.obj),
+            use_glob=use_glob,
+            search_criteria=[item],
+        )
+        if not final_results:
+            print('%s was not found.' % item)
+        else:
+            for res in final_results:
+                results.add(res['name'])
+
+    # Open IPython terminal with RESULTS loaded
 
     logger.debug('Starting load block')
     # retrieve client
     client = get_happi_client_from_config(ctx.obj)
 
     devices = {}
-    names = " ".join(item_names)
+    names = " ".join(results)
     names = names.split()
+
     if len(names) < 1:
         raise click.BadArgumentUsage('No item names given')
-    logger.info(f'Creating shell with devices {item_names}')
+    logger.info(f'Creating shell with devices {results}')
 
     for name in names:
         try:

happi/tests/test_cli.py

@@ -67,8 +67,8 @@ def assert_match_expected(
     """
     logger.debug(result.output)
     assert result.exit_code == expected_return
-    assert not result.exception or (isinstance(result.exception, SystemExit)
-                                    and expected_return != 0)
+    assert not result.exception or (isinstance(result.exception, SystemExit) and
+                                    expected_return != 0)
 
     trimmed_output = trim_split_output(result.output)
     for message, expected in zip(trimmed_output, expected_output):
@@ -91,7 +91,8 @@ def assert_in_expected(
     """
     logger.debug(result.output)
     assert result.exit_code == expected_return
-    assert not result.exception or (isinstance(result.exception, SystemExit) and expected_return != 0)
+    assert not result.exception or (isinstance(
+        result.exception, SystemExit) and expected_return != 0)
     assert expected_inclusion in result.output
 
 
@@ -579,7 +580,7 @@ def test_bad_edit(edit_args: list[str], happi_cfg: str, runner: CliRunner):
     assert bad_edit_result.exit_code == 1
 
 
-def test_load(
+def test_load_one_arg(
     caplog: pytest.LogCaptureFixture,
     client: happi.client.Client,
     happi_cfg: str,
@@ -616,6 +617,69 @@ def test_load(
         assert "Creating shell with devices" in caplog.text
 
 
+def test_load_multiple_args(
+    caplog: pytest.LogCaptureFixture,
+    client: happi.client.Client,
+    happi_cfg: str,
+    runner: CliRunner
+):
+    # try to load the item
+    devices = {}
+    devices['tst_base_pim'] = client.load_device(name='tst_base_pim')
+    devices['tst_base_pim2'] = client.load_device(name='tst_base_pim2')
+    with mock.patch.object(IPython, 'start_ipython') as m:
+        _ = runner.invoke(
+            happi_cli, ['--path', happi_cfg, 'load', 'tst_base_pim', 'tst_base_pim2']
+        )
+        m.assert_called_once_with(argv=['--quick'], user_ns=devices)
+    with caplog.at_level(logging.INFO):
+        assert "Creating shell with devices" in caplog.text
+
+
+def test_load_glob_one_arg(
+    caplog: pytest.LogCaptureFixture,
+    client: happi.client.Client,
+    happi_cfg: str,
+    runner: CliRunner
+):
+    # try to load the item
+    devices = {}
+    devices['tst_base_pim'] = client.load_device(name='tst_base_pim')
+    devices['tst_base_pim2'] = client.load_device(name='tst_base_pim2')
+    devices['tst_minimal'] = client.load_device(name='tst_minimal')
+
+    with mock.patch.object(IPython, 'start_ipython') as m:
+        _ = runner.invoke(
+            happi_cli, ['--path', happi_cfg, 'load', 'tst_*']
+        )
+        m.assert_called_once_with(argv=['--quick'], user_ns=devices)
+
+    with caplog.at_level(logging.INFO):
+        assert "Creating shell with devices" in caplog.text
+
+
+def test_load_glob_multiple_args(
+    caplog: pytest.LogCaptureFixture,
+    client: happi.client.Client,
+    happi_cfg: str,
+    runner: CliRunner
+):
+    # try to load the item
+    devices = {}
+    devices['tst_base_pim'] = client.load_device(name='tst_base_pim')
+    devices['tst_base_pim2'] = client.load_device(name='tst_base_pim2')
+    devices['tst_minimal'] = client.load_device(name='tst_minimal')
+
+    with mock.patch.object(IPython, 'start_ipython') as m:
+        _ = runner.invoke(
+            happi_cli, ['--path', happi_cfg, 'load', 'tst_*', 'device_class=types*']
+        )
+        m.assert_called_once_with(argv=['--quick'], user_ns=devices)
+
+    with caplog.at_level(logging.INFO):
+        assert "Creating shell with devices" in caplog.text
+
+
 @pytest.mark.parametrize("wrapper", ['{{ "tst_base_pim2":{} }}', "[{}]"])
 def test_update(wrapper: str, happi_cfg: str, runner: CliRunner):
     new = """{
@@ -789,7 +853,7 @@ def arg_variants(variants: tuple[tuple[tuple[str]]]):
     """
     for idx, arg_set in enumerate(itertools.product(*variants), 1):
         item = functools.reduce(
-            lambda x, y: x+y,
+            lambda x, y: x + y,
             arg_set,
         )
         summary = f"args{idx}_" + ",".join(item)