Add AsyncConcatenateIterator (#67)

* Add AsyncConcatenateIterator

* Add Async test runners

* Fix recognition of AsyncIterator

Author: 8W9aG

python/cog/__init__.py

@@ -1,5 +1,6 @@
 import coglet
 from coglet.api import (
+    AsyncConcatenateIterator,
     BaseModel,
     BasePredictor,
     CancelationException,
@@ -14,6 +15,7 @@
 __version__ = coglet.__version__
 
 __all__ = [
+    'AsyncConcatenateIterator',
     'BaseModel',
     'BasePredictor',
     'CancelationException',

python/coglet/api.py

@@ -1,7 +1,7 @@
 import pathlib
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Iterator, List, Optional, Type, TypeVar, Union
+from typing import Any, AsyncIterator, Iterator, List, Optional, Type, TypeVar, Union
 
 ########################################
 # Custom encoding
@@ -73,6 +73,11 @@ class ConcatenateIterator(Iterator[_T_co]):
     def __next__(self) -> _T_co: ...
 
 
+class AsyncConcatenateIterator(AsyncIterator[_T_co]):
+    @abstractmethod
+    async def __anext__(self) -> _T_co: ...
+
+
 ########################################
 # Input, Output
 ########################################

python/coglet/inspector.py

@@ -4,7 +4,7 @@
 import typing
 import warnings
 from types import ModuleType
-from typing import Any, Callable, Dict, Iterator, Optional, Type
+from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, Type
 
 from coglet import adt, api
 
@@ -139,13 +139,13 @@ def _output_adt(tpe: type) -> adt.Output:
     origin = typing.get_origin(tpe)
     kind = None
     ft = None
-    if origin is typing.get_origin(Iterator):
+    if origin in {typing.get_origin(Iterator), typing.get_origin(AsyncIterator)}:
         kind = adt.Kind.ITERATOR
         t_args = typing.get_args(tpe)
         assert len(t_args) == 1, 'iterator type must have one type argument'
         ft = adt.FieldType.from_type(t_args[0])
         assert ft.repetition is adt.Repetition.REQUIRED
-    elif origin is api.ConcatenateIterator:
+    elif origin in {api.ConcatenateIterator, api.AsyncConcatenateIterator}:
         kind = adt.Kind.CONCAT_ITERATOR
         t_args = typing.get_args(tpe)
         assert len(t_args) == 1, 'iterator type must have one type argument'

python/tests/runners/async_concat_iterator.py

@@ -0,0 +1,19 @@
+import asyncio
+
+from cog import AsyncConcatenateIterator, BasePredictor
+
+
+class Predictor(BasePredictor):
+    test_inputs = {'i': 3, 's': 'foo'}
+
+    async def predict(self, i: int, s: str) -> AsyncConcatenateIterator[str]:
+        await asyncio.sleep(0.1)
+        print('starting prediction')
+        if i > 0:
+            await asyncio.sleep(0.6)
+        for x in range(i):
+            print(f'prediction in progress {x + 1}/{i}')
+            await asyncio.sleep(0.6)
+            yield f'*{s}-{x}*'
+            await asyncio.sleep(0.6)
+        print('completed prediction')

python/tests/runners/async_iterator.py

@@ -1,13 +1,13 @@
 import asyncio
-from typing import Iterator
+from typing import AsyncIterator
 
 from cog import BasePredictor
 
 
 class Predictor(BasePredictor):
     test_inputs = {'i': 3, 's': 'foo'}
 
-    async def predict(self, i: int, s: str) -> Iterator[str]:
+    async def predict(self, i: int, s: str) -> AsyncIterator[str]:
         await asyncio.sleep(0.1)
         print('starting prediction')
         if i > 0: