From 71b37e78d609d49718fb213ca7895bcf305827a0 Mon Sep 17 00:00:00 2001
From: Pierre Cheynier <p.cheynier@criteo.com>
Date: Mon, 21 Dec 2020 13:30:51 +0100
Subject: [PATCH] Add support for context-managers

By implementing (a)enter/(a)exit internal methods, we bring the support
for context-managers in the base.Consul object.

Context managers are quite convenient, especially for the aio implementation
in which you may need a good control over objects lifecycle to prevent
errors related to already-closed event loops:

This is done by forcing http implementations to define a close method.
Most of the synchronous implementations will use a noop.

```
async with aio.Consul(...) as consul:
  svcs = consul.agent.services()
  # Implicit close of the HTTP client

```
---
 consul/base.py    | 16 ++++++++++++++++
 consul/std.py     |  3 +++
 consul/tornado.py |  3 +++
 consul/twisted.py |  4 ++++
 4 files changed, 26 insertions(+)

diff --git a/consul/base.py b/consul/base.py
index 529e41d..fb35944 100755
--- a/consul/base.py
+++ b/consul/base.py
@@ -277,6 +277,10 @@ def delete(self, callback, path, params=None):
     def post(self, callback, path, params=None, data=''):
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def close(self):
+        raise NotImplementedError
+
 
 class Consul:
     def __init__(
@@ -345,6 +349,18 @@ def __init__(
         self.operator = Consul.Operator(self)
         self.connect = Consul.Connect(self)
 
+    def __enter__(self):
+        return self
+
+    async def __aenter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.http.close()
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.http.close()
+
     class Event:
         """
         The event command provides a mechanism to fire a custom user event to
diff --git a/consul/std.py b/consul/std.py
index 7ecd807..2a23af6 100644
--- a/consul/std.py
+++ b/consul/std.py
@@ -38,6 +38,9 @@ def post(self, callback, path, params=None, data=''):
             self.session.post(uri, data=data, verify=self.verify,
                               cert=self.cert)))
 
+    def close(self):
+        pass
+
 
 class Consul(base.Consul):
     @staticmethod
diff --git a/consul/tornado.py b/consul/tornado.py
index 0a656eb..54ecdcb 100644
--- a/consul/tornado.py
+++ b/consul/tornado.py
@@ -51,6 +51,9 @@ def post(self, callback, path, params=None, data=''):
                                          validate_cert=self.verify)
         return self._request(callback, request)
 
+    def close(self):
+        self.client.close()
+
 
 class Consul(base.Consul):
     @staticmethod
diff --git a/consul/twisted.py b/consul/twisted.py
index 182f88a..0488a16 100644
--- a/consul/twisted.py
+++ b/consul/twisted.py
@@ -121,6 +121,10 @@ def delete(self, callback, path, params=None):
         response = yield self.request(callback, 'delete', uri, params=params)
         returnValue(response)
 
+    @inlineCallbacks
+    def close(self):
+        pass
+
 
 class Consul(base.Consul):
     @staticmethod