From 358976e0b60bce29b9fe8c59807fafbad3e62cbc Mon Sep 17 00:00:00 2001
From: Florimond Manca <florimond.manca@protonmail.com>
Date: Sun, 10 Dec 2023 14:47:54 +0100
Subject: [PATCH] Add anyio 4.x support (#43)

* Add anyio 4.x support

* Update pyproject.toml
---
 .coveragerc_py37           |  4 ++++
 ci/azure-pipelines.yml     |  4 +++-
 pyproject.toml             |  3 ++-
 requirements.txt           |  3 +++
 src/aiometer/_compat.py    | 19 +++++++++++++++++++
 src/aiometer/_impl/amap.py | 39 +++++++++++++++++++++-----------------
 6 files changed, 53 insertions(+), 19 deletions(-)
 create mode 100644 .coveragerc_py37
 create mode 100644 src/aiometer/_compat.py

diff --git a/.coveragerc_py37 b/.coveragerc_py37
new file mode 100644
index 0000000..a424d2f
--- /dev/null
+++ b/.coveragerc_py37
@@ -0,0 +1,4 @@
+[report]
+exclude_lines =
+    pragma: no cover
+    pragma: nopy37
diff --git a/ci/azure-pipelines.yml b/ci/azure-pipelines.yml
index 7dbff8c..e2d57d8 100644
--- a/ci/azure-pipelines.yml
+++ b/ci/azure-pipelines.yml
@@ -30,7 +30,9 @@ stages:
       - template: job--python-test.yml@templates
         parameters:
           jobs:
-            py37: null
+            py37:
+              variables:
+                PYTEST_ADDOPTS: "--cov-config=.coveragerc_py37"
             py311:
               coverage: true
 
diff --git a/pyproject.toml b/pyproject.toml
index e2b4301..953b94a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,8 +24,9 @@ classifiers = [
   "Programming Language :: Python :: 3.11",
 ]
 dependencies = [
-  "anyio~=3.2",
+  "anyio>=3.2,<5",
   "typing-extensions; python_version<'3.8'",
+  "exceptiongroup; python_version<'3.11'",
 ]
 dynamic = ["version", "readme"]
 
diff --git a/requirements.txt b/requirements.txt
index 6e3b886..eff3a20 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,8 @@
 -e .
 
+# Compatibility testing.
+anyio~=3.2; python_version<'3.11'
+
 # Packaging.
 twine
 wheel
diff --git a/src/aiometer/_compat.py b/src/aiometer/_compat.py
new file mode 100644
index 0000000..eb96ab0
--- /dev/null
+++ b/src/aiometer/_compat.py
@@ -0,0 +1,19 @@
+import sys
+from contextlib import contextmanager
+from typing import Generator
+
+if sys.version_info < (3, 11):  # pragma: no cover
+    from exceptiongroup import BaseExceptionGroup
+
+
+@contextmanager
+def collapse_excgroups() -> Generator[None, None, None]:
+    try:
+        yield
+    except BaseException as exc:
+        while (
+            isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1
+        ):  # pragma: nopy37
+            exc = exc.exceptions[0]
+
+        raise exc
diff --git a/src/aiometer/_impl/amap.py b/src/aiometer/_impl/amap.py
index 35c828b..d251a49 100644
--- a/src/aiometer/_impl/amap.py
+++ b/src/aiometer/_impl/amap.py
@@ -14,7 +14,9 @@
 )
 
 import anyio
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
+from .._compat import collapse_excgroups
 from .run_on_each import run_on_each
 from .types import T, U
 
@@ -61,27 +63,30 @@ def amap(
 ) -> AsyncContextManager[AsyncIterable]:
     @asynccontextmanager
     async def _amap() -> AsyncIterator[AsyncIterable]:
-        send_channel, receive_channel = anyio.create_memory_object_stream(
-            max_buffer_size=len(args)
-        )
+        channels: Tuple[
+            MemoryObjectSendStream, MemoryObjectReceiveStream
+        ] = anyio.create_memory_object_stream(max_buffer_size=len(args))
+
+        send_channel, receive_channel = channels
 
         with send_channel, receive_channel:
-            async with anyio.create_task_group() as task_group:
+            with collapse_excgroups():
+                async with anyio.create_task_group() as task_group:
 
-                async def sender() -> None:
-                    # Make any `async for ... in results: ...` terminate.
-                    with send_channel:
-                        await run_on_each(
-                            async_fn,
-                            args,
-                            max_at_once=max_at_once,
-                            max_per_second=max_per_second,
-                            _include_index=_include_index,
-                            _send_to=send_channel,
-                        )
+                    async def sender() -> None:
+                        # Make any `async for ... in results: ...` terminate.
+                        with send_channel:
+                            await run_on_each(
+                                async_fn,
+                                args,
+                                max_at_once=max_at_once,
+                                max_per_second=max_per_second,
+                                _include_index=_include_index,
+                                _send_to=send_channel,
+                            )
 
-                task_group.start_soon(sender)
+                    task_group.start_soon(sender)
 
-                yield receive_channel
+                    yield receive_channel
 
     return _amap()