Include AsyncIterator import for both clients and servers (#264)

Co-authored-by: Robin Lambertz <github@roblab.la>
danielgtaylor · Nov 5, 2021 · c424b6f · c424b6f
1 parent 421fdba
commit c424b6f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/betterproto/grpc/grpclib_server.py b/src/betterproto/grpc/grpclib_server.py
@@ -1,5 +1,5 @@
 from abc import ABC
-from collections import AsyncIterable
+from collections.abc import AsyncIterable
 from typing import Callable, Any, Dict
 
 import grpclib

diff --git a/src/betterproto/plugin/models.py b/src/betterproto/plugin/models.py
@@ -653,7 +653,9 @@ def __post_init__(self) -> None:
             self.output_file.typing_imports.add("AsyncIterable")
             self.output_file.typing_imports.add("Iterable")
             self.output_file.typing_imports.add("Union")
-        if self.server_streaming:
+
+        # Required by both client and server
+        if self.client_streaming or self.server_streaming:
             self.output_file.typing_imports.add("AsyncIterator")
 
         super().__post_init__()  # check for unset fields