From 3b91105aea40644aba020e0059a8c1315f00757c Mon Sep 17 00:00:00 2001
From: Felix Stupp <felix.stupp@outlook.com>
Date: Fri, 3 Apr 2020 10:06:17 +0000
Subject: [PATCH] Connection: Added support for with statement

---
 tinytinypy/main.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tinytinypy/main.py b/tinytinypy/main.py
index 73e5a86..e5b12bd 100644
--- a/tinytinypy/main.py
+++ b/tinytinypy/main.py
@@ -106,7 +106,7 @@ class Connection:
         self._endpoint = endpoint
         self._sid = None
         self.__conn = protos[proto](host=host)
-        atexit.register(lambda: self.__conn.close())
+        atexit.register(lambda: self.close())
 
     def __raiseError(self, op, info):
         if type(info) == 'dict':
@@ -119,6 +119,17 @@ class Connection:
             err = info
         raise Exception(f"API Call {op} failed: {info['error']}")
 
+    # For supporting with statement
+    def __enter__(self):
+        return self
+    def __exit__(self, ex_type, value, tb):
+        self.close()
+        return True if ex_type is None else False
+
+    def close(self):
+        self.logout()
+        self.__conn.close()
+
     def _get(self, op, sendSid=True, **parameters):
         headers = {
                 'Content-Type': 'application/json',