From f3f047c69383f160a9891c109a7ea0324221fae8 Mon Sep 17 00:00:00 2001
From: Vincent Michel <vxgmichel@gmail.com>
Date: Fri, 3 May 2024 20:06:59 +0200
Subject: [PATCH] Fix covariance issue

---
 aiostream/aiter_utils.py | 2 +-
 aiostream/core.py        | 2 +-
 tests/test_advanced.py   | 5 +----
 3 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/aiostream/aiter_utils.py b/aiostream/aiter_utils.py
index 0ef74fe..df5213d 100644
--- a/aiostream/aiter_utils.py
+++ b/aiostream/aiter_utils.py
@@ -105,7 +105,7 @@ def assert_async_iterator(obj: object) -> None:
 
 # Async iterator context
 
-T = TypeVar("T")
+T = TypeVar("T", covariant=True)
 Self = TypeVar("Self", bound="AsyncIteratorContext[Any]")
 
 
diff --git a/aiostream/core.py b/aiostream/core.py
index 9bf4fdc..0be5e7d 100644
--- a/aiostream/core.py
+++ b/aiostream/core.py
@@ -38,7 +38,7 @@ class StreamEmpty(Exception):
 
 # Helpers
 
-T = TypeVar("T")
+T = TypeVar("T", covariant=True)
 X = TypeVar("X")
 A = TypeVar("A", contravariant=True)
 P = ParamSpec("P")
diff --git a/tests/test_advanced.py b/tests/test_advanced.py
index 4df9ca8..e085d59 100644
--- a/tests/test_advanced.py
+++ b/tests/test_advanced.py
@@ -13,10 +13,7 @@ def target2(x: int, *_) -> Stream[int]:
         return stream.range(x, x + 4, interval=1)
 
     def target3(x: int, *_) -> Stream[int]:
-        # TODO: fix covariance issue
-        return (
-            stream.range(0, 3, interval=1) if x else stream.throw(ZeroDivisionError)
-        )  # type: ignore
+        return stream.range(0, 3, interval=1) if x else stream.throw(ZeroDivisionError)
 
     # Concurrent run
     with assert_cleanup() as loop: