diff --git a/package/AUTHORS b/package/AUTHORS index 245e6bd7345..e4d426d89c0 100644 --- a/package/AUTHORS +++ b/package/AUTHORS @@ -214,6 +214,7 @@ Chronological list of authors - Alexander Schlaich - Josh Vermaas - Xiaoxu Ruan + - Egor Marin External code ------------- diff --git a/package/CHANGELOG b/package/CHANGELOG index 2d489056e09..4e073688b81 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -37,6 +37,7 @@ Fixes (Issue #3336) Enhancements + * Add `progressbar_kwargs` parameter to `AnalysisBase.run` method, allowing to modify description, position etc of tqdm progressbars. * Add a nojump transformation, which unwraps trajectories so that particle paths are continuous. (Issue #3703, PR #4031) * Added AtomGroup TopologyAttr to calculate gyration moments (Issue #3904, diff --git a/package/MDAnalysis/analysis/base.py b/package/MDAnalysis/analysis/base.py index 266ae68a02e..4e3638cc18c 100644 --- a/package/MDAnalysis/analysis/base.py +++ b/package/MDAnalysis/analysis/base.py @@ -390,7 +390,7 @@ def _conclude(self): pass # pylint: disable=unnecessary-pass def run(self, start=None, stop=None, step=None, frames=None, - verbose=None): + verbose=None, *, progressbar_kwargs={}): """Perform the calculation Parameters @@ -411,12 +411,19 @@ def run(self, start=None, stop=None, step=None, frames=None, verbose : bool, optional Turn on verbosity + + progressbar_kwargs : dict, optional + ProgressBar keywords with custom parameters regarding progress bar position, etc; + see :class:`MDAnalysis.lib.log.ProgressBar` for full list. .. versionchanged:: 2.2.0 Added ability to analyze arbitrary frames by passing a list of frame indices in the `frames` keyword argument. + .. versionchanged:: 2.5.0 + Add `progressbar_kwargs` parameter, + allowing to modify description, position etc of tqdm progressbars """ logger.info("Choosing frames to analyze") # if verbose unchanged, use class default @@ -429,9 +436,11 @@ def run(self, start=None, stop=None, step=None, frames=None, self._prepare() logger.info("Starting analysis loop over %d trajectory frames", self.n_frames) + for i, ts in enumerate(ProgressBar( self._sliced_trajectory, - verbose=verbose)): + verbose=verbose, + **progressbar_kwargs)): self._frame_index = i self._ts = ts self.frames[i] = ts.frame diff --git a/testsuite/MDAnalysisTests/analysis/test_base.py b/testsuite/MDAnalysisTests/analysis/test_base.py index 09b7b158f0a..e26cee6912b 100644 --- a/testsuite/MDAnalysisTests/analysis/test_base.py +++ b/testsuite/MDAnalysisTests/analysis/test_base.py @@ -274,6 +274,12 @@ def test_verbose_progressbar_run(u, capsys): actual = err.strip().split('\r')[-1] assert actual[:24] == expected[:24] +def test_verbose_progressbar_run_with_kwargs(u, capsys): + an = FrameAnalysis(u.trajectory).run(verbose=True, progressbar_kwargs={'desc':'custom'}) + out, err = capsys.readouterr() + expected = u'custom: 100%|██████████| 98/98 [00:00<00:00, 8799.49it/s]' + actual = err.strip().split('\r')[-1] + assert actual[:30] == expected[:30] def test_incomplete_defined_analysis(u): with pytest.raises(NotImplementedError):