Make servers and stubs context managers

Servers and stubs were context managers in the Alpha API; they may not need to be in the Beta API but it's easy enough to do, eases migration, and probably helps some use cases. For now the stub is given empty __enter__ and __exit__ methods; we can always come back and implement the actual work of context management in a later change.
lenmoyouziJiangjun · Sep 8, 2015 · c824eb4 · c824eb4
1 parent 0269c9c
commit c824eb4
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/src/python/grpcio/grpc/beta/_server.py b/src/python/grpcio/grpc/beta/_server.py
@@ -86,13 +86,13 @@ def add_secure_port(self, address, server_credentials):
     return self._grpc_link.add_port(
         address, server_credentials._intermediary_low_credentials)  # pylint: disable=protected-access
 
-  def start(self):
+  def _start(self):
     self._grpc_link.join_link(self._end_link)
     self._end_link.join_link(self._grpc_link)
     self._grpc_link.start()
     self._end_link.start()
 
-  def stop(self, grace):
+  def _stop(self, grace):
     stop_event = threading.Event()
     if 0 < grace:
       disassembly_thread = threading.Thread(
@@ -105,6 +105,20 @@ def stop(self, grace):
       _disassemble(self._grpc_link, self._end_link, self._pool, stop_event, 0)
       return stop_event
 
+  def start(self):
+    self._start()
+
+  def stop(self, grace):
+    return self._stop(grace)
+
+  def __enter__(self):
+    self._start()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    self._stop(0).wait()
+    return False
+
 
 def server(
     implementations, multi_implementation, request_deserializers,

diff --git a/src/python/grpcio/grpc/beta/_stub.py b/src/python/grpcio/grpc/beta/_stub.py
@@ -49,6 +49,12 @@ def __init__(self, delegate, on_deletion):
   def __getattr__(self, attr):
     return getattr(self._delegate, attr)
 
+  def __enter__(self):
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    return False
+
   def __del__(self):
     self._on_deletion()