Merge pull request #3279 from nathanielmanistaatgoogle/context-management

Make servers and stubs context managers

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()