Repository: cogent3/cogent3
Pull request #1255 (2022-12-29 11:22:24)

ENH: progress bar on apps


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index 50d0a61ba..4e70a828a 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -1288,13 +1288,8 @@ def _source_wrapped(self, value: source_proxy) -> source_proxy:
     return value
 
 
-def _as_completed(
-    self,
-    dstore,
-    parallel=False,
-    par_kw=None,
-    # show_progress?
-) -> Generator:
+@UI.display_wrap
+def _as_completed(self, dstore, parallel=False, par_kw=None, **kwargs) -> Generator:
     """invokes self composable function on the provided data store
 
     Parameters
@@ -1309,10 +1304,10 @@ def _as_completed(
         member of dstore.
     par_kw
         dict of values for configuring parallel execution.
-
-    Returns
-    -------
-    Result of the process as a list.
+    kwargs
+        setting a show_progress boolean keyword value here
+        affects progress display code, other arguments are passed to
+        the cogent3.util.progress_bar.display_wrap decorator
 
     Notes
     -----
@@ -1320,6 +1315,7 @@ def _as_completed(
     aggregates results. If run in serial, results are returned in the
     same order as provided.
     """
+    ui = kwargs.pop("ui")
     app = (
         self.input._source_wrapped if self.app_type is WRITER else self._source_wrapped
     )
@@ -1335,7 +1331,8 @@ def _as_completed(
         to_do = PAR.as_completed(app, mapped, **par_kw)
     else:
         to_do = map(app, mapped)
-    return to_do
+
+    return ui.series(to_do, count=len(mapped), **kwargs)
 
 
 def is_composable(obj):
@@ -1377,15 +1374,18 @@ def _apply_to(
     cleanup : bool
         after copying of log files into the data store, it is deleted
         from the original location
+    show_progress : bool
+        controls progress bar display
 
     Returns
     -------
-    Result of the process as a list
+    The output data store instance
 
     Notes
     -----
     This is an append only function, meaning that if a member already exists
     in self.data_store for an input, it is skipped.
+
     If run in parallel, this instance spawns workers and aggregates results.
     """
     if not self.input:
@@ -1415,7 +1415,9 @@ def _apply_to(
     logger.log_versions(["cogent3"])
 
     inputs = _proxy_input(inputs.values())
-    for result in self.as_completed(inputs, parallel=parallel, par_kw=par_kw):
+    for result in self.as_completed(
+        inputs, parallel=parallel, par_kw=par_kw, show_progress=show_progress
+    ):
         member = self.main(
             data=result.obj, identifier=id_from_source(result.source)
         )  # writers must return DataMember
diff --git a/src/cogent3/app/io_new.py b/src/cogent3/app/io_new.py
index b02baab6e..581d5c7b2 100644
--- a/src/cogent3/app/io_new.py
+++ b/src/cogent3/app/io_new.py
@@ -39,7 +39,7 @@ from .data_store_new import (
     load_record_from_json,
     make_record_for_json,
 )
-from .sqlite_data_store import DataStoreSqlite
+from .sqlite_data_store import _MEMORY, DataStoreSqlite
 from .typing import (
     AlignedSeqsType,
     IdentifierType,
@@ -141,13 +141,15 @@ def open_data_store(
 
     kwargs = dict(limit=limit, mode=mode, suffix=suffix)
     base_path = Path(base_path)
-    base_path = base_path.expanduser().absolute()
+    base_path = (
+        base_path if base_path.name == ":memory:" else base_path.expanduser().absolute()
+    )
     if base_path.suffix == ".tinydb":
         kwargs["suffix"] = "json"
         kwargs.pop("mode")
 
-    if base_path.suffix == ".sqlitedb":
-        ds_suffix = base_path.suffix
+    if base_path.suffix == ".sqlitedb" or base_path.name == _MEMORY:
+        ds_suffix = ".sqlitedb"
         kwargs.pop("suffix")
     elif zipfile.is_zipfile(base_path):
         ds_suffix = ".zip"
@@ -157,6 +159,9 @@ def open_data_store(
     else:
         ds_suffix = None
 
+    if base_path.name == _MEMORY and mode is READONLY:
+        raise NotImplementedError("in memory readonly sqlitedb")
+
     if ds_suffix is None and suffix is None:
         raise ValueError("a suffix is required if using a directory data store")
 


Test patch:
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index dbe3684e7..e768b6ad4 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -436,6 +436,32 @@ def test_apply_to_not_partially_done(tmp_dir):
     assert len(out_dstore) == num_records
 
 
+@pytest.mark.xfail(reason="passes except when run in full test suite")
+@pytest.mark.parametrize("show", (True, False))
+def test_as_completed_progress(full_dstore, show):
+    # this contextlib trap works only if test is run separately
+    # but clearly there's some effect on stderr being
+    # re-piped by other tests?
+    # the capsys fixture does not seem to work at all, but is likely how to
+    # do this robustly
+    import io
+
+    from contextlib import redirect_stderr
+
+    loader = get_app("load_unaligned", format="fasta", moltype="dna")
+    min_length = get_app("min_length", 300)
+    app = loader + min_length
+    with redirect_stderr(io.StringIO()) as f:
+        list(app.as_completed(full_dstore.completed, show_progress=show))
+
+    result = f.getvalue().splitlines()
+    if show:
+        assert len(result) > 0
+        assert "100%" in result[-1]
+    else:
+        assert len(result) == 0
+
+
 def test_err_result():
     """excercise creation of NotCompletedResult"""
     result = NotCompleted("SKIP", "this", "some obj")
diff --git a/tests/test_app/test_sqlite_data_store.py b/tests/test_app/test_sqlite_data_store.py
index 6ee418612..a90d0f047 100644
--- a/tests/test_app/test_sqlite_data_store.py
+++ b/tests/test_app/test_sqlite_data_store.py
@@ -6,6 +6,7 @@ import pytest
 
 from scitrack import get_text_hexdigest
 
+from cogent3 import open_data_store
 from cogent3.app.composable import NotCompleted
 from cogent3.app.data_store_new import (
     APPEND,
@@ -541,3 +542,9 @@ def test_validate_missing_md5(md5_none):
     assert t["Num md5sum missing", "Value"] == 9
     for c in ("correct", "incorrect"):
         assert t[f"Num md5sum {c}", "Value"] == 0
+
+
+def test_open_data_store_sqlitedb_err():
+    # cannot create an in-mmemory db to read only
+    with pytest.raises(NotImplementedError):
+        open_data_store(":memory:", mode="r")