Merge pull request PyMySQL#26 from PyMySQL/feature/cursor-context-manager

methane · methane · commit a42fe113bf7b · 2015-02-25T11:38:05.000+09:00
Add context manager interface to cursor
diff --git a/MySQLdb/cursors.py b/MySQLdb/cursors.py
@@ -85,13 +85,22 @@ def __init__(self, connection):
         
     def close(self):
         """Close the cursor. No further queries will be possible."""
-        if self.connection is None or self.connection() is None:
-            return
-        while self.nextset():
-            pass
-        self.connection = None
-        self.errorhandler = None
-        self._result = None
+        try:
+            if self.connection is None or self.connection() is None:
+                return
+            while self.nextset():
+                pass
+        finally:
+            self.connection = None
+            self.errorhandler = None
+            self._result = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        del exc_info
+        self.close()
 
     def _check_executed(self):
         if not self._executed:
