From 15b63b7190cca08a414b56433856ea521aef7731 Mon Sep 17 00:00:00 2001
From: Tyler Goodlet <jgbt@protonmail.com>
Date: Fri, 7 May 2021 11:52:08 -0400
Subject: [PATCH] Fix typing

---
 tractor/_portal.py    |  2 +-
 tractor/_streaming.py | 20 ++++++++++++++++----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/tractor/_portal.py b/tractor/_portal.py
index e13792b..8944489 100644
--- a/tractor/_portal.py
+++ b/tractor/_portal.py
@@ -340,7 +340,7 @@ class Portal:
         self,
         func: Callable,
         **kwargs,
-    ) -> Context:
+    ) -> AsyncGenerator[Tuple[Context, Any], None]:
         """Open an inter-actor task context.
 
         This is a synchronous API which allows for deterministic
diff --git a/tractor/_streaming.py b/tractor/_streaming.py
index 10be2c6..9bc32e3 100644
--- a/tractor/_streaming.py
+++ b/tractor/_streaming.py
@@ -1,7 +1,15 @@
+"""
+Message stream types and APIs.
+
+"""
 import inspect
 from contextlib import contextmanager, asynccontextmanager
 from dataclasses import dataclass
-from typing import Any, Iterator, Optional, Callable
+from typing import (
+    Any, Iterator, Optional, Callable,
+    AsyncGenerator,
+)
+
 import warnings
 
 import trio
@@ -288,7 +296,7 @@ class Context:
     async def open_stream(
         self,
         shield: bool = False,
-    ) -> MsgStream:
+    ) -> AsyncGenerator[MsgStream, None]:
         # TODO
 
         actor = current_actor()
@@ -339,7 +347,9 @@ def stream(func: Callable) -> Callable:
 
     """
     # annotate
-    func._tractor_stream_function = True
+    # TODO: apply whatever solution ``mypy`` ends up picking for this:
+    # https://github.com/python/mypy/issues/2087#issuecomment-769266912
+    func._tractor_stream_function = True  # type: ignore
 
     sig = inspect.signature(func)
     params = sig.parameters
@@ -369,7 +379,9 @@ def context(func: Callable) -> Callable:
 
     """
     # annotate
-    func._tractor_context_function = True
+    # TODO: apply whatever solution ``mypy`` ends up picking for this:
+    # https://github.com/python/mypy/issues/2087#issuecomment-769266912
+    func._tractor_context_function = True  # type: ignore
 
     sig = inspect.signature(func)
     params = sig.parameters