Support context managers, and improve request methods for rest client (#577)

Author: bl4ckst0ne

atlassian/rest_client.py

@@ -54,6 +54,12 @@ def __init__(self, url, username=None, password=None, timeout=60, api_root='rest
         elif cookies is not None:
             self._session.cookies.update(cookies)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        self.close()
+
     def _create_basic_session(self, username, password):
         self._session.auth = (username, password)
 
@@ -89,6 +95,17 @@ def _update_header(self, key, value):
         """
         self._session.headers.update({key: value})
 
+    @staticmethod
+    def _response_handler(response):
+        try:
+            return response.json()
+        except ValueError:
+            log.debug('Received response with no content')
+            return None
+        except Exception as e:
+            log.error(e)
+            return None
+
     def log_curl_debug(self, method, url, data=None, headers=None, level=logging.DEBUG):
         """
 
@@ -117,6 +134,9 @@ def url_joiner(url, path, trailing=None):
             url_link += '/'
         return url_link
 
+    def close(self):
+        return self._session.close()
+
     def request(self, method='GET', path='/', data=None, json=None, flags=None, params=None, headers=None,
                 files=None, trailing=None):
         """
@@ -199,28 +219,15 @@ def post(self, path, data=None, json=None, headers=None, files=None, params=None
                                 trailing=trailing)
         if self.advanced_mode:
             return response
-        try:
-            return response.json()
-        except ValueError:
-            log.debug('Received response with no content')
-            return None
-        except Exception as e:
-            log.error(e)
-            return None
+        return self._response_handler(response)
 
     def put(self, path, data=None, headers=None, files=None, trailing=None, params=None):
         response = self.request('PUT', path=path, data=data, headers=headers, files=files, params=params,
                                 trailing=trailing)
         if self.advanced_mode:
             return response
-        try:
-            return response.json()
-        except ValueError:
-            log.debug('Received response with no content')
-            return None
-        except Exception as e:
-            log.error(e)
-            return None
+        return self._response_handler(response)
+
 
     def delete(self, path, data=None, headers=None, params=None, trailing=None):
         """
@@ -232,11 +239,4 @@ def delete(self, path, data=None, headers=None, params=None, trailing=None):
         response = self.request('DELETE', path=path, data=data, headers=headers, params=params, trailing=trailing)
         if self.advanced_mode:
             return response
-        try:
-            return response.json()
-        except ValueError:
-            log.debug('Received response with no content')
-            return None
-        except Exception as e:
-            log.error(e)
-            return None
+        return self._response_handler(response)