Merge branch 'coroutines' into devel

offensity · Jul 16, 2020 · cb60ea9 · cb60ea9
2 parents e7b2e0c + 4f4e108
commit cb60ea9
Show file tree

Hide file tree

Showing 8 changed files with 362 additions and 35 deletions.
diff --git a/.meta/.readme.rst b/.meta/.readme.rst
@@ -376,7 +376,7 @@ Returns
       def set_description(self, desc=None, refresh=True):
           """{DOC_tqdm.tqdm.set_description}"""
 
-      def set_postfix(self, ordered_dict=None, refresh=True, **kwargs):
+      def set_postfix(self, ordered_dict=None, refresh=True, **tqdm_kwargs):
           """{DOC_tqdm.tqdm.set_postfix}"""
 
       @classmethod
@@ -390,29 +390,45 @@ Returns
       def display(self, msg=None, pos=None):
           """{DOC_tqdm.tqdm.display}"""
 
-    def trange(*args, **kwargs):
+      @classmethod
+      @contextmanager
+      def wrapattr(cls, stream, method, total=None, bytes=True, **tqdm_kwargs):
+          """{DOC_tqdm.tqdm.wrapattr}"""
+
+      @classmethod
+      def pandas(cls, *targs, **tqdm_kwargs):
+          """Registers the current `tqdm` class with `pandas`."""
+
+    def trange(*args, **tqdm_kwargs):
         """
-        A shortcut for tqdm(xrange(*args), **kwargs).
-        On Python3+ range is used instead of xrange.
+        A shortcut for `tqdm(xrange(*args), **tqdm_kwargs)`.
+        On Python3+, `range` is used instead of `xrange`.
         """
 
-    class tqdm.gui.tqdm(tqdm.tqdm):
-        """Experimental GUI version"""
+    class tqdm.notebook.tqdm(tqdm.tqdm):
+        """Experimental IPython/Jupyter Notebook widget."""
 
-    def tqdm.gui.trange(*args, **kwargs):
-        """Experimental GUI version of trange"""
+    def tqdm.notebook.trange(*args, **tqdm_kwargs):
+        """Experimental IPython/Jupyter Notebook widget version of `trange`."""
 
-    class tqdm.notebook.tqdm(tqdm.tqdm):
-        """Experimental IPython/Jupyter Notebook widget"""
+    class tqdm.auto.tqdm(tqdm.tqdm):
+        """Automatically chooses beween `tqdm.notebook` and `tqdm.tqdm`."""
 
-    def tqdm.notebook.trange(*args, **kwargs):
-        """Experimental IPython/Jupyter Notebook widget version of trange"""
+    class tqdm.asyncio.tqdm(tqdm.tqdm):
+      """Asynchronous version."""
+      @classmethod
+      def as_completed(cls, fs, *, loop=None, timeout=None, total=None,
+                       **tqdm_kwargs):
+          """Wrapper for `asyncio.as_completed`."""
+
+    class tqdm.gui.tqdm(tqdm.tqdm):
+        """Experimental GUI version."""
 
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
-        """`keras` callback for epoch and batch progress"""
+        """`keras` callback for epoch and batch progress."""
 
     def tqdm.contrib.tenumerate(iterable, start=0, total=None,
-                                tqdm_class=tqdm.auto.tqdm, **kwargs):
+                                tqdm_class=tqdm.auto.tqdm, **tqdm_kwargs):
         """Equivalent of `numpy.ndenumerate` or builtin `enumerate`."""
 
     def tqdm.contrib.tzip(iter1, *iter2plus, **tqdm_kwargs):
@@ -428,13 +444,12 @@ The ``tqdm.contrib`` package also contains experimental modules:
 
 - ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
 - ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
+- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
+- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
 - ``tqdm.contrib.bells``: Automagically enables all optional features
 
   * ``auto``, ``pandas``, ``discord``, ``telegram``
 
-- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
-- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
-
 Examples and Advanced Usage
 ---------------------------
 

diff --git a/README.rst b/README.rst
@@ -543,7 +543,7 @@ Returns
               Forces refresh [default: True].
           """
 
-      def set_postfix(self, ordered_dict=None, refresh=True, **kwargs):
+      def set_postfix(self, ordered_dict=None, refresh=True, **tqdm_kwargs):
           """
           Set/modify postfix (additional stats)
           with automatic formatting based on datatype.
@@ -578,29 +578,55 @@ Returns
             (default: ``abs(self.pos)``).
           """
 
-    def trange(*args, **kwargs):
+      @classmethod
+      @contextmanager
+      def wrapattr(cls, stream, method, total=None, bytes=True, **tqdm_kwargs):
+          """
+          stream  : file-like object.
+          method  : str, "read" or "write". The result of ``read()`` and
+              the first argument of ``write()`` should have a ``len()``.
+
+          >>> with tqdm.wrapattr(file_obj, "read", total=file_obj.size) as fobj:
+          ...     while True:
+          ...         chunk = fobj.read(chunk_size)
+          ...         if not chunk:
+          ...             break
+          """
+
+      @classmethod
+      def pandas(cls, *targs, **tqdm_kwargs):
+          """Registers the current `tqdm` class with `pandas`."""
+
+    def trange(*args, **tqdm_kwargs):
         """
-        A shortcut for tqdm(xrange(*args), **kwargs).
-        On Python3+ range is used instead of xrange.
+        A shortcut for `tqdm(xrange(*args), **tqdm_kwargs)`.
+        On Python3+, `range` is used instead of `xrange`.
         """
 
-    class tqdm.gui.tqdm(tqdm.tqdm):
-        """Experimental GUI version"""
+    class tqdm.notebook.tqdm(tqdm.tqdm):
+        """Experimental IPython/Jupyter Notebook widget."""
 
-    def tqdm.gui.trange(*args, **kwargs):
-        """Experimental GUI version of trange"""
+    def tqdm.notebook.trange(*args, **tqdm_kwargs):
+        """Experimental IPython/Jupyter Notebook widget version of `trange`."""
 
-    class tqdm.notebook.tqdm(tqdm.tqdm):
-        """Experimental IPython/Jupyter Notebook widget"""
+    class tqdm.auto.tqdm(tqdm.tqdm):
+        """Automatically chooses beween `tqdm.notebook` and `tqdm.tqdm`."""
 
-    def tqdm.notebook.trange(*args, **kwargs):
-        """Experimental IPython/Jupyter Notebook widget version of trange"""
+    class tqdm.asyncio.tqdm(tqdm.tqdm):
+      """Asynchronous version."""
+      @classmethod
+      def as_completed(cls, fs, *, loop=None, timeout=None, total=None,
+                       **tqdm_kwargs):
+          """Wrapper for `asyncio.as_completed`."""
+
+    class tqdm.gui.tqdm(tqdm.tqdm):
+        """Experimental GUI version."""
 
     class tqdm.keras.TqdmCallback(keras.callbacks.Callback):
-        """`keras` callback for epoch and batch progress"""
+        """`keras` callback for epoch and batch progress."""
 
     def tqdm.contrib.tenumerate(iterable, start=0, total=None,
-                                tqdm_class=tqdm.auto.tqdm, **kwargs):
+                                tqdm_class=tqdm.auto.tqdm, **tqdm_kwargs):
         """Equivalent of `numpy.ndenumerate` or builtin `enumerate`."""
 
     def tqdm.contrib.tzip(iter1, *iter2plus, **tqdm_kwargs):
@@ -616,13 +642,12 @@ The ``tqdm.contrib`` package also contains experimental modules:
 
 - ``tqdm.contrib.itertools``: Thin wrappers around ``itertools``
 - ``tqdm.contrib.concurrent``: Thin wrappers around ``concurrent.futures``
+- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
+- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
 - ``tqdm.contrib.bells``: Automagically enables all optional features
 
   * ``auto``, ``pandas``, ``discord``, ``telegram``
 
-- ``tqdm.contrib.discord``: Posts to `Discord <https://discord.com/>`__ bots
-- ``tqdm.contrib.telegram``: Posts to `Telegram <https://telegram.org/>`__ bots
-
 Examples and Advanced Usage
 ---------------------------
 

diff --git a/examples/async_coroutines.py b/examples/async_coroutines.py
@@ -0,0 +1,37 @@
+"""
+Asynchronous examples using `asyncio`, `async` and `await` on `python>=3.7`.
+"""
+import asyncio
+from tqdm.asyncio import tqdm, trange
+
+
+def count(start=0, step=1):
+    i = start
+    while True:
+        new_start = yield i
+        if new_start is None:
+            i += step
+        else:
+            i = new_start
+
+
+async def main():
+    N = int(1e6)
+    async for row in tqdm(trange(N, desc="inner"), desc="outer"):
+        if row >= N:
+            break
+    with tqdm(count(), desc="coroutine", total=N + 2) as pbar:
+        async for row in pbar:
+            if row == N:
+                pbar.send(-10)
+            elif row < 0:
+                assert row == -9
+                break
+    # should be under 10 seconds
+    for i in tqdm.as_completed(list(map(asyncio.sleep, [1] * 10)),
+                               desc="as_completed"):
+        await i
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/examples/coroutine_pipe.py b/examples/coroutine_pipe.py
@@ -0,0 +1,68 @@
+"""
+Inserting `tqdm` as a "pipe" in a chain of coroutines.
+Not to be confused with `asyncio.coroutine`.
+"""
+from functools import wraps
+from tqdm.auto import tqdm
+
+
+def autonext(func):
+    @wraps(func)
+    def inner(*args, **kwargs):
+        res = func(*args, **kwargs)
+        next(res)
+        return res
+    return inner
+
+
+@autonext
+def tqdm_pipe(target, **tqdm_kwargs):
+    """
+    Coroutine chain pipe `send()`ing to `target`.
+
+    This:
+    >>> r = receiver()
+    >>> p = producer(r)
+    >>> next(r)
+    >>> next(p)
+
+    Becomes:
+    >>> r = receiver()
+    >>> t = tqdm.pipe(r)
+    >>> p = producer(t)
+    >>> next(r)
+    >>> next(p)
+    """
+    with tqdm(**tqdm_kwargs) as pbar:
+        while True:
+            obj = (yield)
+            target.send(obj)
+            pbar.update()
+
+
+def source(target):
+    for i in ["foo", "bar", "baz", "pythonista", "python", "py"]:
+        target.send(i)
+    target.close()
+
+
+@autonext
+def grep(pattern, target):
+    while True:
+        line = (yield)
+        if pattern in line:
+            target.send(line)
+
+
+@autonext
+def sink():
+    while True:
+        line = (yield)
+        tqdm.write(line)
+
+
+if __name__ == "__main__":
+    source(
+        tqdm_pipe(
+            grep('python',
+                 sink())))
diff --git a/tqdm/asyncio.py b/tqdm/asyncio.py
@@ -0,0 +1,72 @@
+"""
+Asynchronous progressbar decorator for iterators.
+Includes a default `range` iterator printing to `stderr`.
+
+Usage:
+>>> from tqdm.asyncio import trange, tqdm
+>>> async for i in trange(10):
+...     ...
+"""
+from .auto import tqdm as tqdm_auto
+import asyncio
+__author__ = {"github.com/": ["casperdcl"]}
+__all__ = ['tqdm_asyncio', 'tarange', 'tqdm', 'trange']
+
+
+class tqdm_asyncio(tqdm_auto):
+    def __init__(self, iterable=None, *args, **kwargs):
+        super(tqdm_asyncio, self).__init__(iterable, *args, **kwargs)
+        self.iterable_awaitable = False
+        if iterable is not None:
+            if hasattr(iterable, "__anext__"):
+                self.iterable_next = iterable.__anext__
+                self.iterable_awaitable = True
+            elif hasattr(iterable, "__next__"):
+                self.iterable_next = iterable.__next__
+            else:
+                self.iterable_iterator = iter(iterable)
+                self.iterable_next = self.iterable_iterator.__next__
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            if self.iterable_awaitable:
+                res = await self.iterable_next()
+            else:
+                res = self.iterable_next()
+            self.update()
+            return res
+        except StopIteration:
+            self.close()
+            raise StopAsyncIteration
+        except:
+            self.close()
+            raise
+
+    def send(self, *args, **kwargs):
+        return self.iterable.send(*args, **kwargs)
+
+    @classmethod
+    def as_completed(cls, fs, *, loop=None, timeout=None, total=None,
+                     **tqdm_kwargs):
+        """
+        Wrapper for `asyncio.as_completed`.
+        """
+        if total is None:
+            total = len(fs)
+        yield from cls(asyncio.as_completed(fs, loop=loop, timeout=timeout),
+                       total=total, **tqdm_kwargs)
+
+
+def tarange(*args, **kwargs):
+    """
+    A shortcut for `tqdm.asyncio.tqdm(range(*args), **kwargs)`.
+    """
+    return tqdm_asyncio(range(*args), **kwargs)
+
+
+# Aliases
+tqdm = tqdm_asyncio
+trange = tarange
diff --git a/tqdm/std.py b/tqdm/std.py
@@ -630,7 +630,7 @@ def get_lock(cls):
     @classmethod
     def pandas(cls, **tqdm_kwargs):
         """
-        Registers the given `tqdm` class with
+        Registers the current `tqdm` class with
             pandas.core.
             ( frame.DataFrame
             | series.Series