diff --git a/MySQLdb/connections.py b/MySQLdb/connections.py index 27a04770..020b4a99 100644 --- 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 index fa4692ef..c5cacbec 100644 --- 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)