From ee5f8d21e88b59fe89abfcbc8762f66b2d648ade Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <tqdm@caspersci.uk.to>
Date: Fri, 20 Nov 2020 23:27:22 +0000
Subject: [PATCH] fix nested `asyncio`, document async `break` hazard

- fixes #1074
---
 .meta/.readme.rst     | 25 +++++++++++++++++++++++++
 README.rst            | 25 +++++++++++++++++++++++++
 tests/py37_asyncio.py | 21 +++++++++++----------
 tqdm/asyncio.py       | 20 ++++++++++++++++++++
 4 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/.meta/.readme.rst b/.meta/.readme.rst
index b9852c43..5f9e0f4f 100644
--- a/.meta/.readme.rst
+++ b/.meta/.readme.rst
@@ -744,6 +744,31 @@ custom callback take advantage of this, simply use the return value of
                 external_callback(**self.format_dict)
             return displayed
 
+``asyncio``
+~~~~~~~~~~~
+
+Note that ``break`` isn't currently caught by asynchronous iterators.
+This means that ``tqdm`` cannot clean up after itself in this case:
+
+.. code:: python
+
+    from tqdm.asyncio import tqdm
+
+    async for i in tqdm(range(9)):
+        if i == 2:
+            break
+
+Instead, either call ``pbar.close()`` manually or use the context manager syntax:
+
+.. code:: python
+
+    from tqdm.asyncio import tqdm
+
+    with tqdm(range(9)) as pbar:
+        async for i in pbar:
+            if i == 2:
+                break
+
 Pandas Integration
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/README.rst b/README.rst
index f54b42f5..0284e60b 100644
--- a/README.rst
+++ b/README.rst
@@ -961,6 +961,31 @@ custom callback take advantage of this, simply use the return value of
                 external_callback(**self.format_dict)
             return displayed
 
+``asyncio``
+~~~~~~~~~~~
+
+Note that ``break`` isn't currently caught by asynchronous iterators.
+This means that ``tqdm`` cannot clean up after itself in this case:
+
+.. code:: python
+
+    from tqdm.asyncio import tqdm
+
+    async for i in tqdm(range(9)):
+        if i == 2:
+            break
+
+Instead, either call ``pbar.close()`` manually or use the context manager syntax:
+
+.. code:: python
+
+    from tqdm.asyncio import tqdm
+
+    with tqdm(range(9)) as pbar:
+        async for i in pbar:
+            if i == 2:
+                break
+
 Pandas Integration
 ~~~~~~~~~~~~~~~~~~
 
diff --git a/tests/py37_asyncio.py b/tests/py37_asyncio.py
index 1f079645..cd9389d7 100644
--- a/tests/py37_asyncio.py
+++ b/tests/py37_asyncio.py
@@ -14,8 +14,8 @@ as_completed = partial(tqdm_asyncio.as_completed, miniters=0, mininterval=0)
 
 def with_setup_sync(func):
     @wraps(func)
-    def inner():
-        return asyncio.run(func())
+    def inner(*args, **kwargs):
+        return asyncio.run(func(*args, **kwargs))
     return inner
 
 
@@ -35,20 +35,21 @@ async def acount(*args, **kwargs):
 
 
 @with_setup_sync
-async def test_generators():
+async def test_generators(capsys):
     """Test asyncio generators"""
-    with closing(StringIO()) as our_file:
-        async for i in tqdm(count(), desc="counter", file=our_file):
+    with tqdm(count(), desc="counter") as pbar:
+        async for i in pbar:
             if i >= 8:
                 break
-        assert '9it' in our_file.getvalue()
-        our_file.seek(0)
-        our_file.truncate()
+    _, err = capsys.readouterr()
+    assert '9it' in err
 
-        async for i in tqdm(acount(), desc="async_counter", file=our_file):
+    with tqdm(acount(), desc="async_counter") as pbar:
+        async for i in pbar:
             if i >= 8:
                 break
-        assert '9it' in our_file.getvalue()
+    _, err = capsys.readouterr()
+    assert '9it' in err
 
 
 @with_setup_sync
diff --git a/tqdm/asyncio.py b/tqdm/asyncio.py
index d7f08a99..e3e2d276 100644
--- a/tqdm/asyncio.py
+++ b/tqdm/asyncio.py
@@ -62,6 +62,26 @@ class tqdm_asyncio(std_tqdm):
         yield from cls(asyncio.as_completed(fs, loop=loop, timeout=timeout),
                        total=total, **tqdm_kwargs)
 
+    def __new__(cls, *args, **kwargs):
+        """
+        Workaround for mixed-class same-stream nested progressbars.
+        See [#509](https://github.com/tqdm/tqdm/issues/509)
+        """
+        with cls.get_lock():
+            try:
+                cls._instances = std_tqdm._instances
+            except AttributeError:
+                pass
+        instance = super(tqdm_asyncio, cls).__new__(cls, *args, **kwargs)
+        with cls.get_lock():
+            try:
+                # `std_tqdm` may have been changed so update
+                cls._instances.update(std_tqdm._instances)
+            except AttributeError:
+                pass
+            std_tqdm._instances = cls._instances
+        return instance
+
 
 def tarange(*args, **kwargs):
     """