Add Context Manager Interface to Connection. (PyMySQL#413)

methane · web-flow · commit 643220972e01 · 2019-11-29T21:15:17.000+09:00
Fixes PyMySQL#400.
diff --git a/MySQLdb/connections.py b/MySQLdb/connections.py
@@ -207,6 +207,12 @@ def unicode_literal(u, dummy=None):
                 self.autocommit(autocommit)
         self.messages = []
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def autocommit(self, on):
         on = bool(on)
         if self.get_autocommit() != on:
diff --git a/tests/test_MySQLdb_nonstandard.py b/tests/test_MySQLdb_nonstandard.py
@@ -99,3 +99,8 @@ def test_client_flag(self):
 
     def test_fileno(self):
         self.assertGreaterEqual(self.conn.fileno(), 0)
+
+    def test_context_manager(self):
+        with connection_factory() as conn:
+            self.assertFalse(conn.closed)
+        self.assertTrue(conn.closed)
