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.

2 files changed, 22 insertions, 2 deletions

diff --git a/src/python/grpcio/grpc/beta/_server.py b/src/python/grpcio/grpc/beta/_server.py
index daa42c475a..05b954d186 100644
--- a/src/python/grpcio/grpc/beta/_server.py
+++ b/src/python/grpcio/grpc/beta/_server.py
@@ -86,13 +86,13 @@ class Server(interfaces.Server):
     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 @@ class Server(interfaces.Server):
       _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
index cfbecb852b..11dab889cd 100644
--- a/src/python/grpcio/grpc/beta/_stub.py
+++ b/src/python/grpcio/grpc/beta/_stub.py
@@ -49,6 +49,12 @@ class _AutoIntermediary(object):
   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()