fix(when): fix type overloads of then_enter_with on mypy v0.941 (#124)

Author: mcous

decoy/__init__.py

@@ -1,16 +1,30 @@
 """Decoy stubbing and spying library."""
 from warnings import warn
-from typing import Any, Callable, Generic, Optional, Union, cast, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    Optional,
+    Union,
+    cast,
+    overload,
+)
 
 from . import errors, matchers, warnings
-from .context_managers import (
-    AsyncContextManager,
-    ContextManager,
-    GeneratorContextManager,
-)
 from .core import DecoyCore, StubCore, PropCore
 from .types import ClassT, ContextValueT, FuncT, ReturnT
 
+# TODO(mc, 2022-03-14): drop support for Python 3.6 in Decoy v2
+# Python 3.6 does not have async generator context managers
+if TYPE_CHECKING:
+    from .context_managers import (
+        ContextManager,
+        AsyncContextManager,
+        GeneratorContextManager,
+        AsyncGeneratorContextManager,
+    )
+
 # ensure decoy does not pollute pytest tracebacks
 __tracebackhide__ = True
 
@@ -294,11 +308,19 @@ def then_enter_with(
     ) -> None:
         ...
 
+    @overload
+    def then_enter_with(
+        self: "Stub[AsyncGeneratorContextManager[ContextValueT]]",
+        value: ContextValueT,
+    ) -> None:
+        ...
+
     def then_enter_with(
         self: Union[
-            "Stub[GeneratorContextManager[ContextValueT]]",
             "Stub[ContextManager[ContextValueT]]",
+            "Stub[GeneratorContextManager[ContextValueT]]",
             "Stub[AsyncContextManager[ContextValueT]]",
+            "Stub[AsyncGeneratorContextManager[ContextValueT]]",
         ],
         value: ContextValueT,
     ) -> None:

decoy/context_managers.py

@@ -1,8 +1,20 @@
 """Wrappers around contextlib types and fallbacks."""
-import contextlib
-from typing import Any, AsyncContextManager, ContextManager, Generic, TypeVar
-
-GeneratorContextManager = contextlib._GeneratorContextManager
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    ContextManager,
+    Generic,
+    TypeVar,
+)
+
+# TODO(mc, 2022-03-14): drop support for Python 3.6 in Decoy v2
+# Python 3.6 does not have async generator context managers
+if TYPE_CHECKING:
+    from contextlib import (  # type: ignore[attr-defined]
+        _GeneratorContextManager as GeneratorContextManager,
+        _AsyncGeneratorContextManager as AsyncGeneratorContextManager,
+    )
 
 _EnterT = TypeVar("_EnterT")
 
@@ -40,7 +52,8 @@ async def __aexit__(self, *args: Any, **kwargs: Any) -> Any:
 
 __all__ = [
     "AsyncContextManager",
-    "GeneratorContextManager",
+    "AsyncGeneratorContextManager",
     "ContextManager",
+    "GeneratorContextManager",
     "ContextWrapper",
 ]