Merge branch 'wraps' into devel

KOLANICH-libs · Jan 24, 2020 · ad04dc8 · ad04dc8
2 parents af40b04 + ad87480
commit ad04dc8
Show file tree

Hide file tree

Showing 14 changed files with 381 additions and 25 deletions.
diff --git a/.meta/.readme.rst b/.meta/.readme.rst
@@ -280,13 +280,18 @@ of a neat one-line progress bar.
     either normal-width unicode characters being incorrectly displayed as
     "wide", or some unicode characters not rendering.
 
-- Wrapping enumerated iterables: use ``enumerate(tqdm(...))`` instead of
-  ``tqdm(enumerate(...))``. The same applies to ``numpy.ndenumerate``.
-  This is because enumerate functions tend to hide the length of iterables.
-  ``tqdm`` does not.
-- Wrapping zipped iterables has similar issues due to internal optimisations.
-  ``tqdm(zip(a, b))`` should be replaced with ``zip(tqdm(a), b)`` or even
-  ``zip(tqdm(a), tqdm(b))``.
+- Wrapping generators:
+
+  * Generator wrapper functions tend to hide the length of iterables.
+    ``tqdm`` does not.
+  * Replace ``tqdm(enumerate(...))`` with ``enumerate(tqdm(...))`` or
+    ``tqdm(enumerate(x), total=len(x), ...)``.
+    The same applies to ``numpy.ndenumerate``.
+  * Replace ``tqdm(zip(a, b))`` with ``zip(tqdm(a), b)`` or even
+    ``zip(tqdm(a), tqdm(b))``.
+  * The same applies to ``itertools``.
+  * Some useful convenience functions can be found under ``tqdm.contrib``.
+
 - `Hanging pipes in python2 <https://github.com/tqdm/tqdm/issues/359>`__:
   when using ``tqdm`` on the CLI, you may need to use Python 3.5+ for correct
   buffering.
@@ -381,6 +386,21 @@ Returns
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
         """`keras` callback for epoch and batch progress"""
 
+    def tqdm.contrib.tenumerate(iterable, start=0, total=None,
+                                tqdm_class=tqdm.auto.tqdm, **kwargs):
+        """Equivalent of `numpy.ndenumerate` or builtin `enumerate`."""
+
+    def tqdm.contrib.tzip(iter1, *iter2plus, **tqdm_kwargs):
+        """Equivalent of builtin `zip`."""
+
+    def tqdm.contrib.tmap(function, *sequences, **tqdm_kwargs):
+        """Equivalent of builtin `map`."""
+
+The ``tqdm.contrib`` package also contains experimental modules:
+
+- ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
+- ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
+
 Examples and Advanced Usage
 ---------------------------
 

diff --git a/Makefile b/Makefile
@@ -118,10 +118,12 @@ prebuildclean:
 coverclean:
 	@+python -c "import os; os.remove('.coverage') if os.path.exists('.coverage') else None"
 	@+python -c "import shutil; shutil.rmtree('tqdm/__pycache__', True)"
+	@+python -c "import shutil; shutil.rmtree('tqdm/contrib/__pycache__', True)"
 	@+python -c "import shutil; shutil.rmtree('tqdm/tests/__pycache__', True)"
 clean:
 	@+python -c "import os, glob; [os.remove(i) for i in glob.glob('*.py[co]')]"
 	@+python -c "import os, glob; [os.remove(i) for i in glob.glob('tqdm/*.py[co]')]"
+	@+python -c "import os, glob; [os.remove(i) for i in glob.glob('tqdm/contrib/*.py[co]')]"
 	@+python -c "import os, glob; [os.remove(i) for i in glob.glob('tqdm/tests/*.py[co]')]"
 	@+python -c "import os, glob; [os.remove(i) for i in glob.glob('tqdm/examples/*.py[co]')]"
 toxclean:

diff --git a/README.rst b/README.rst
@@ -280,13 +280,18 @@ of a neat one-line progress bar.
     either normal-width unicode characters being incorrectly displayed as
     "wide", or some unicode characters not rendering.
 
-- Wrapping enumerated iterables: use ``enumerate(tqdm(...))`` instead of
-  ``tqdm(enumerate(...))``. The same applies to ``numpy.ndenumerate``.
-  This is because enumerate functions tend to hide the length of iterables.
-  ``tqdm`` does not.
-- Wrapping zipped iterables has similar issues due to internal optimisations.
-  ``tqdm(zip(a, b))`` should be replaced with ``zip(tqdm(a), b)`` or even
-  ``zip(tqdm(a), tqdm(b))``.
+- Wrapping generators:
+
+  * Generator wrapper functions tend to hide the length of iterables.
+    ``tqdm`` does not.
+  * Replace ``tqdm(enumerate(...))`` with ``enumerate(tqdm(...))`` or
+    ``tqdm(enumerate(x), total=len(x), ...)``.
+    The same applies to ``numpy.ndenumerate``.
+  * Replace ``tqdm(zip(a, b))`` with ``zip(tqdm(a), b)`` or even
+    ``zip(tqdm(a), tqdm(b))``.
+  * The same applies to ``itertools``.
+  * Some useful convenience functions can be found under ``tqdm.contrib``.
+
 - `Hanging pipes in python2 <https://github.com/tqdm/tqdm/issues/359>`__:
   when using ``tqdm`` on the CLI, you may need to use Python 3.5+ for correct
   buffering.
@@ -563,6 +568,21 @@ Returns
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
         """`keras` callback for epoch and batch progress"""
 
+    def tqdm.contrib.tenumerate(iterable, start=0, total=None,
+                                tqdm_class=tqdm.auto.tqdm, **kwargs):
+        """Equivalent of `numpy.ndenumerate` or builtin `enumerate`."""
+
+    def tqdm.contrib.tzip(iter1, *iter2plus, **tqdm_kwargs):
+        """Equivalent of builtin `zip`."""
+
+    def tqdm.contrib.tmap(function, *sequences, **tqdm_kwargs):
+        """Equivalent of builtin `map`."""
+
+The ``tqdm.contrib`` package also contains experimental modules:
+
+- ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
+- ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
+
 Examples and Advanced Usage
 ---------------------------
 

diff --git a/examples/include_no_requirements.py b/examples/include_no_requirements.py
@@ -1,6 +1,6 @@
-# How to import tqdm without enforcing it as a dependency
+# How to import tqdm in any frontend without enforcing it as a dependency
 try:
-    from tqdm import tqdm
+    from tqdm.auto import tqdm
 except ImportError:
 
     def tqdm(*args, **kwargs):

diff --git a/examples/pandas_progress_apply.py b/examples/pandas_progress_apply.py
@@ -1,6 +1,6 @@
 import pandas as pd
 import numpy as np
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 df = pd.DataFrame(np.random.randint(0, 100, (100000, 6)))
 

diff --git a/examples/parallel_bars.py b/examples/parallel_bars.py
@@ -1,6 +1,7 @@
 from __future__ import print_function
 from time import sleep
-from tqdm import tqdm, trange
+from tqdm.auto import tqdm, trange
+from tqdm.contrib.concurrent import process_map, thread_map
 from random import random
 from multiprocessing import Pool, freeze_support
 from concurrent.futures import ThreadPoolExecutor
@@ -12,11 +13,12 @@
 PY2 = sys.version_info[:1] <= (2,)
 
 
-def progresser(n, auto_position=True, write_safe=False, blocking=True):
+def progresser(n, auto_position=True, write_safe=False, blocking=True,
+               progress=False):
     interval = random() * 0.002 / (NUM_SUBITERS - n + 2)
     total = 5000
     text = "#{}, est. {:<04.2}s".format(n, interval * total)
-    for _ in trange(total, desc=text,
+    for _ in trange(total, desc=text, disable=not progress,
                     lock_args=None if blocking else (False,),
                     position=None if auto_position else n):
         sleep(interval)
@@ -26,20 +28,27 @@ def progresser(n, auto_position=True, write_safe=False, blocking=True):
         # we think we know about other bars (currently only py3 threading)
         if n == 6:
             tqdm.write("n == 6 completed")
+    return n + 1
 
 
 if __name__ == '__main__':
     freeze_support()  # for Windows support
     L = list(range(NUM_SUBITERS))[::-1]
 
+    print("Simple thread mapping")
+    thread_map(partial(progresser, write_safe=not PY2), L, max_workers=4)
+
+    print("Simple process mapping")
+    process_map(partial(progresser), L, max_workers=4)
+
     print("Manual nesting")
     for i in trange(16, desc="1"):
         for _ in trange(16, desc="2 @ %d" % i, leave=i % 2):
             sleep(0.01)
 
     print("Multi-processing")
     p = Pool(initializer=tqdm.set_lock, initargs=(tqdm.get_lock(),))
-    p.map(progresser, L)
+    p.map(partial(progresser, progress=True), L)
 
     # unfortunately need ncols
     # to print spaces over leftover multi-processing bars (#796)
@@ -50,6 +59,5 @@ def progresser(n, auto_position=True, write_safe=False, blocking=True):
     # explicitly set just threading lock for nonblocking progress
     tqdm.set_lock(RLock())
     with ThreadPoolExecutor() as p:
-        progresser_thread = partial(
-            progresser, write_safe=not PY2, blocking=False)
-        p.map(progresser_thread, L)
+        p.map(partial(progresser, progress=True, write_safe=not PY2,
+                      blocking=False), L)
diff --git a/examples/tqdm_wget.py b/examples/tqdm_wget.py
@@ -21,7 +21,7 @@
 
 import urllib
 from os import devnull
-from tqdm import tqdm
+from tqdm.auto import tqdm
 from docopt import docopt
 
 

diff --git a/examples/wrapping_generators.py b/examples/wrapping_generators.py
@@ -0,0 +1,14 @@
+from tqdm.contrib import tenumerate, tzip, tmap
+import numpy as np
+
+for _ in tenumerate(range(int(1e6)), desc="builtin enumerate"):
+    pass
+
+for _ in tenumerate(np.random.random((999, 999)), desc="numpy.ndenumerate"):
+    pass
+
+for _ in tzip(np.arange(1e6), np.arange(1e6) + 1, desc="builtin zip"):
+    pass
+
+mapped = tmap(lambda x: x + 1, np.arange(1e6), desc="builtin map")
+assert (np.arange(1e6) + 1 == list(mapped)).all()
diff --git a/tqdm/contrib/__init__.py b/tqdm/contrib/__init__.py
@@ -1,5 +1,16 @@
+"""
+Thin wrappers around common functions.
+
+Subpackages contain potentially unstable extensions.
+"""
 from tqdm import tqdm
+from tqdm.auto import tqdm as tqdm_auto
 from tqdm.utils import ObjectWrapper
+from copy import deepcopy
+import functools
+import sys
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['tenumerate', 'tzip', 'tmap']
 
 
 class DummyTqdmFile(ObjectWrapper):
@@ -8,3 +19,62 @@ def write(self, x, nolock=False):
         # Avoid print() second call (useless \n)
         if len(x.rstrip()) > 0:
             tqdm.write(x, file=self._wrapped, nolock=nolock)
+
+
+def tenumerate(iterable, start=0, total=None, tqdm_class=tqdm_auto,
+               **tqdm_kwargs):
+    """
+    Equivalent of `numpy.ndenumerate` or builtin `enumerate`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    try:
+        import numpy as np
+    except ImportError:
+        pass
+    else:
+        if isinstance(iterable, np.ndarray):
+            return tqdm_class(np.ndenumerate(iterable),
+                              total=total or iterable.size, **tqdm_kwargs)
+    return enumerate(tqdm_class(iterable, **tqdm_kwargs), start)
+
+
+def _tzip(iter1, *iter2plus, **tqdm_kwargs):
+    """
+    Equivalent of builtin `zip`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    kwargs = deepcopy(tqdm_kwargs)
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    for i in zip(tqdm_class(iter1, **tqdm_kwargs), *iter2plus):
+        yield i
+
+
+def _tmap(function, *sequences, **tqdm_kwargs):
+    """
+    Equivalent of builtin `map`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    for i in _tzip(*sequences, **tqdm_kwargs):
+        yield function(*i)
+
+
+if sys.version_info[:1] < (3,):
+    @functools.wraps(_tzip)
+    def tzip(*args, **kwargs):
+        return list(_tzip(*args, **kwargs))
+
+    @functools.wraps(_tmap)
+    def tmap(*args, **kwargs):
+        return list(_tmap(*args, **kwargs))
+else:
+    tzip = _tzip
+    tmap = _tmap
diff --git a/tqdm/contrib/concurrent.py b/tqdm/contrib/concurrent.py
@@ -0,0 +1,64 @@
+"""
+Thin wrappers around `concurrent.futures`.
+"""
+from __future__ import absolute_import
+from tqdm.auto import tqdm as tqdm_auto
+from copy import deepcopy
+try:
+    from os import cpu_count
+except ImportError:
+    try:
+        from multiprocessing import cpu_count
+    except ImportError:
+        def cpu_count():
+            return 4
+import sys
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['thread_map', 'process_map']
+
+
+def _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs):
+    """
+    Implementation of `thread_map` and `process_map`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    kwargs = deepcopy(tqdm_kwargs)
+    kwargs.setdefault("total", len(iterables[0]))
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4))
+    pool_kwargs = dict(max_workers=max_workers)
+    if sys.version_info[:2] >= (3, 7):
+        # share lock in case workers are already using `tqdm`
+        pool_kwargs.update(
+            initializer=tqdm_class.set_lock, initargs=(tqdm_class.get_lock(),))
+    with PoolExecutor(**pool_kwargs) as ex:
+        return list(tqdm_class(ex.map(fn, *iterables), **kwargs))
+
+
+def thread_map(fn, *iterables, **tqdm_kwargs):
+    """
+    Equivalent of `list(map(fn, *iterables))`
+    driven by `concurrent.futures.ThreadPoolExecutor`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    from concurrent.futures import ThreadPoolExecutor
+    return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs)
+
+
+def process_map(fn, *iterables, **tqdm_kwargs):
+    """
+    Equivalent of `list(map(fn, *iterables))`
+    driven by `concurrent.futures.ProcessPoolExecutor`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    from concurrent.futures import ProcessPoolExecutor
+    return _executor_map(ProcessPoolExecutor, fn, *iterables, **tqdm_kwargs)
diff --git a/tqdm/contrib/itertools.py b/tqdm/contrib/itertools.py
@@ -0,0 +1,34 @@
+"""
+Thin wrappers around `itertools`.
+"""
+from __future__ import absolute_import
+from tqdm.auto import tqdm as tqdm_auto
+from copy import deepcopy
+import itertools
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['product']
+
+
+def product(*iterables, **tqdm_kwargs):
+    """
+    Equivalent of `itertools.product`.
+
+    Parameters
+    ----------
+    tqdm_class  : [default: tqdm.auto.tqdm].
+    """
+    kwargs = deepcopy(tqdm_kwargs)
+    tqdm_class = kwargs.pop("tqdm_class", tqdm_auto)
+    try:
+        lens = list(map(len, iterables))
+    except TypeError:
+        total = None
+    else:
+        total = 1
+        for i in lens:
+            total *= i
+        kwargs.setdefault("total", total)
+    with tqdm_class(**kwargs) as t:
+        for i in itertools.product(*iterables):
+            yield i
+            t.update()