Progress bar follow-up (#110)

* doc fixes and tweaks

* rename progress.py -> monitoring.py

Also don't expose ProgressBar directly in xsimlab main namespace.

* test hook functions

* black

* update release notes

doc/api.rst

@@ -170,13 +170,13 @@ listed below. These are defined in ``xsimlab.validators``.
 
 .. _`attrs' validators`: https://www.attrs.org/en/stable/examples.html#validators
 
-Runtime monitoring
-==================
+Model runtime monitoring
+========================
 
 .. currentmodule:: xsimlab
 .. autosummary::
    :toctree: _api_generated/
 
+   monitoring.ProgressBar
    runtime_hook
    RuntimeHook
-   ProgressBar

doc/monitor.rst

@@ -41,42 +41,53 @@ The following imports are necessary for the examples below.
 Progress bar
 ------------
 
-:class:`~xsimlab.ProgressBar` is based on the `Tqdm`_ package and allows to track
-the progress of simulation runs in ``xarray-simlab``.
-It can be used as a context manager around simulation calls:
+:class:`~xsimlab.monitoring.ProgressBar` is based on the `Tqdm`_ package and
+allows to track the progress of simulation runs in ``xarray-simlab``. It can be
+used as a context manager around simulation calls:
 
-.. _Tqdm: https://github.com/tqdm/tqdm/
+.. _Tqdm: https://tqdm.github.io
 
-.. ipython::
+.. ipython:: python
 
-   In [2]: with xs.ProgressBar():
-      ...:     out_ds = in_ds.xsimlab.run(model=model2)
+   from xsimlab.monitoring import ProgressBar
 
-Alternatively, you can pass the progress bar via the ``hooks`` argument or use the
-``register`` method (for more information, refer to the :ref:`custom_runtime_hooks` subsection)
+.. ipython:: python
+   :suppress:
 
-``ProgressBar`` and the underlying Tqdm is built to work with different Python
-interfaces. Use the optional argument ``frontend`` according to your
-development environment.
+   from progress_bar_hack import ProgressBarHack as ProgressBar
 
-- ``auto``: (default) Automatically detects environment.
-- ``console``: When Python is run from the command line.
-- ``gui``: Tqdm provides a gui version. According to the developers, this is
-  still an experimental feature.
-- ``notebook``: For use in a IPython/Jupyter notebook.
+.. ipython:: python
 
-Additionally, you can customize the built-in progress bar, by supplying a
-keyworded argument list to ``ProgressBar``, e.g.:
+   with ProgressBar():
+       out_ds = in_ds.xsimlab.run(model=model2)
 
-.. ipython::
+Alternatively, you can pass the progress bar via the ``hooks`` argument of
+``Dataset.xsimlab.run()`` or you can use the ``register`` method (for more
+information, refer to the :ref:`custom_runtime_hooks` subsection).
 
-   In [4]: with xs.ProgressBar(bar_format="{r_bar}"):
-      ...:     out_ds = in_ds.xsimlab.run(model=model2)
+``ProgressBar`` and the underlying Tqdm tool are built to work with different
+Python front-ends. Use the optional argument ``frontend`` depending on your
+environment:
+
+- ``auto``: automatically selects the front-end (default)
+- ``console``: renders the progress bar as text
+- ``gui``: progress rich rendering (experimental), which needs matplotlib_ to be
+  installed
+- ``notebook``: for use within IPython/Jupyter notebooks, which needs
+  ipywidgets_ to be installed
+
+.. _matplotlib: https://matplotlib.org/
+.. _ipywidgets: https://ipywidgets.readthedocs.io/en/stable/
+
+Additionally, you can customize the built-in progress bar by supplying
+keyword arguments list to ``ProgressBar``, e.g.:
+
+.. ipython:: python
 
-For a full list of customization options, refer to the `Tqdm documentation`_
+   with ProgressBar(bar_format="{desc}|{bar}{r_bar}"):
+       out_ds = in_ds.xsimlab.run(model=model2)
 
-Note: The ``total`` argument cannot be changed to ensure best performance and
-functionality.
+For a full list of customization options, refer to the `Tqdm documentation`_.
 
 .. _Tqdm documentation: https://tqdm.github.io
 

doc/scripts/progress_bar_hack.py

@@ -0,0 +1,26 @@
+# A hack around ProgressBar (monkey patch) so that it renders
+# nicely in docs
+import io
+
+import xsimlab
+from xsimlab import runtime_hook
+from xsimlab.monitoring import ProgressBar as _ProgressBar
+
+
+class ProgressBarHack(_ProgressBar):
+    """Redirects progress bar outputs to a variable, and
+    only display the rendered string (last line) at the end
+    the simulation.
+
+    """
+
+    def __init__(self, **kwargs):
+        super(ProgressBarHack, self).__init__(**kwargs)
+
+        self.pbar_output = io.StringIO()
+        self.tqdm_kwargs.update({"file": self.pbar_output})
+
+    @runtime_hook("finalize", trigger="post")
+    def close_bar(self, model, context, state):
+        super(ProgressBarHack, self).close_bar(model, context, state)
+        print(self.pbar_output.getvalue().strip().split("\r")[-1])

doc/whats_new.rst

@@ -71,8 +71,8 @@ Enhancements
 - Added some useful properties and methods to the ``xarray.Dataset.xsimlab``
   extension (:issue:`103`).
 - Save model inputs/outputs using zarr (:issue:`102`).
-- Added :class:`~xsimlab.progress.ProgressBar` to track simulation progress
-  (:issue:`104`).
+- Added :class:`~xsimlab.monitoring.ProgressBar` to track simulation progress
+  (:issue:`104`, :issue:`110`).
 
 Bug fixes
 ~~~~~~~~~

xsimlab/__init__.py

@@ -13,9 +13,9 @@
     runtime,
     variable_info,
 )
-from .progress import ProgressBar
 from .variable import variable, index, on_demand, foreign, group
 from .xr_accessor import SimlabAccessor, create_setup
+from . import monitoring
 
 from ._version import get_versions
 

xsimlab/progress.py → xsimlab/monitoring.py

@@ -1,42 +1,51 @@
 from xsimlab.hook import RuntimeHook, runtime_hook
 
 
+__all__ = ("ProgressBar",)
+
+
 class ProgressBar(RuntimeHook):
     """
     Progress bar implementation using the tqdm package.
 
-    Parameters
-    ----------
-    frontend : {"auto", "console", "gui", "notebook"}, optional
-        Selects a frontend for displaying the progress bar. By default ("auto"),
-        the frontend is chosen by guessing in which environment the simulation
-        is run. The "console" frontend displays an ascii progress bar, while the
-        "gui" frontend is based on matplotlib and the "notebook" frontend is based
-        on ipywidgets.
-    **kwargs : dict, optional
-        Arbitrary keyword arguments for progress bar customization.
-
     Examples
     --------
-    :class:`ProgressBar` takes full advantage of :class:`RuntimeHook`.
+    ProgressBar takes full advantage of :class:`RuntimeHook`.
 
-    Call it as part of :func:`run`:
-    >>> out_ds = in_ds.xsimlab.run(model=model, hooks=[xs.ProgressBar()])
+    Call it as part of :meth:`xarray.Dataset.xsimlab.run`:
 
-    In a context manager using the `with` statement`:
-    >>> with xs.ProgressBar():
+    >>> from xsimlab.monitoring import ProgressBar
+    >>> out_ds = in_ds.xsimlab.run(model=model, hooks=[ProgressBar()])
+
+    In a context manager using the ``with`` statement:
+
+    >>> with ProgressBar():
     ...    out_ds = in_ds.xsimlab.run(model=model)
 
-    Globally with `register` method:
-    >>> pbar = xs.ProgressBar()
+    Globally with ``register`` method:
+
+    >>> pbar = ProgressBar()
     >>> pbar.register()
     >>> out_ds = in_ds.xsimlab.run(model=model)
     >>> pbar.unregister()
 
-    For additional customization, see: https://tqdm.github.io/docs/tqdm/
     """
 
     def __init__(self, frontend="auto", **kwargs):
+        """
+        Parameters
+        ----------
+        frontend : {"auto", "console", "gui", "notebook"}, optional
+            Selects a frontend for displaying the progress bar. By default ("auto"),
+            the frontend is chosen by guessing in which environment the simulation
+            is run. The "console" frontend displays an ascii progress bar, while the
+            "gui" frontend is based on matplotlib and the "notebook" frontend is based
+            on ipywidgets.
+        **kwargs : dict, optional
+            Arbitrary keyword arguments for progress bar customization.
+            See https://tqdm.github.io/docs/tqdm/.
+
+        """
         if frontend == "auto":
             from tqdm.auto import tqdm
         elif frontend == "console":
@@ -47,7 +56,9 @@ def __init__(self, frontend="auto", **kwargs):
             from tqdm.notebook import tqdm
         else:
             raise ValueError(
-                f"Frontend argument {frontend!r} not supported. Please select one of the following: {', '.join(['auto', 'console', 'gui', 'notebook'])}"
+                f"Frontend argument {frontend!r} not supported. "
+                "Please select one of the following: "
+                ", ".join(["auto", "console", "gui", "notebook"])
             )
 
         self.custom_description = False
@@ -71,7 +82,7 @@ def update_init(self, mode, context, state):
         self.pbar_model.update(1)
 
     @runtime_hook("run_step", trigger="post")
-    def update_runstep(self, mode, context, state):
+    def update_run_step(self, model, context, state):
         if not self.custom_description:
             self.pbar_model.set_description_str(
                 f"run step {context['step']}/{context['nsteps']}"

xsimlab/tests/test_monitoring.py

@@ -0,0 +1,92 @@
+import importlib
+
+import pytest
+
+from ..monitoring import ProgressBar
+from . import has_tqdm
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+@pytest.mark.parametrize(
+    "frontend,tqdm_module",
+    [
+        ("auto", "tqdm"),  # assume tests are run in a terminal evironment
+        ("console", "tqdm"),
+        ("gui", "tqdm.gui"),
+        ("notebook", "tqdm.notebook"),
+    ],
+)
+def test_progress_bar_init(frontend, tqdm_module):
+    pbar = ProgressBar(frontend=frontend)
+    tqdm = importlib.import_module(tqdm_module)
+
+    assert pbar.tqdm is tqdm.tqdm
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+@pytest.mark.parametrize("kw", [{}, {"bar_format": "{bar}"}])
+def test_progress_bar_init_kwargs(kw):
+    pbar = ProgressBar(**kw)
+
+    assert "bar_format" in pbar.tqdm_kwargs
+
+    if "bar_format" in kw:
+        assert pbar.tqdm_kwargs["bar_format"] == kw["bar_format"]
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_init_error(in_dataset, model):
+    with pytest.raises(ValueError, match=r".*not supported.*"):
+        ProgressBar(frontend="invalid_frontend")
+
+
+@pytest.mark.parametrize("kw", [{}, {"desc": "custom description"}])
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_init_bar(kw):
+    pbar = ProgressBar(**kw)
+    pbar.init_bar(None, {"nsteps": 10}, {})
+
+    assert pbar.pbar_model.format_dict["total"] == 12
+    if kw:
+        assert pbar.pbar_model.format_dict["prefix"] == "custom description"
+    else:
+        assert pbar.pbar_model.format_dict["prefix"] == "initialize"
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_update_init():
+    pbar = ProgressBar()
+    pbar.init_bar(None, {"nsteps": 10}, {})
+    pbar.update_init(None, {}, {})
+
+    assert pbar.pbar_model.format_dict["n"] == 1
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_update_run_step():
+    pbar = ProgressBar()
+    pbar.init_bar(None, {"nsteps": 10}, {})
+    pbar.update_init(None, {}, {})
+    pbar.update_run_step(None, {"nsteps": 10, "step": 1}, {})
+
+    assert pbar.pbar_model.format_dict["n"] == 2
+    assert pbar.pbar_model.format_dict["prefix"] == "run step 1/10"
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_update_finalize():
+    pbar = ProgressBar()
+    pbar.init_bar(None, {"nsteps": 10}, {})
+    pbar.update_finalize(None, {}, {})
+
+    assert pbar.pbar_model.format_dict["prefix"] == "finalize"
+
+
+@pytest.mark.skipif(not has_tqdm, reason="requires tqdm")
+def test_progress_bar_close_bar():
+    pbar = ProgressBar()
+    pbar.init_bar(None, {"nsteps": 10}, {})
+    pbar.close_bar(None, {}, {})
+
+    assert pbar.pbar_model.format_dict["n"] == 1
+    assert pbar.pbar_model.format_dict["prefix"].startswith("Simulation finished")