Adding upstream version 1.2+20240521.upstream/1.2+20240521upstream

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 33 insertions, 0 deletions

diff --git a/src/aristaproto/grpc/grpclib_server.py b/src/aristaproto/grpc/grpclib_server.py
new file mode 100644
index 0000000..3e28031
--- /dev/null
+++ b/src/aristaproto/grpc/grpclib_server.py
@@ -0,0 +1,33 @@
+from abc import ABC
+from collections.abc import AsyncIterable
+from typing import (
+    Any,
+    Callable,
+    Dict,
+)
+
+import grpclib
+import grpclib.server
+
+
+class ServiceBase(ABC):
+    """
+    Base class for async gRPC servers.
+    """
+
+    async def _call_rpc_handler_server_stream(
+        self,
+        handler: Callable,
+        stream: grpclib.server.Stream,
+        request: Any,
+    ) -> None:
+        response_iter = handler(request)
+        # check if response is actually an AsyncIterator
+        # this might be false if the method just returns without
+        # yielding at least once
+        # in that case, we just interpret it as an empty iterator
+        if isinstance(response_iter, AsyncIterable):
+            async for response_message in response_iter:
+                await stream.send_message(response_message)
+        else:
+            response_iter.close()