Context manager support for pure Python Gaviota

niklasf · Oct 16, 2015 · 83cf049 · 83cf049
1 parent b55af43
commit 83cf049
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/chess/gaviota.py b/chess/gaviota.py
@@ -2071,6 +2071,12 @@ def egtb_block_park(self, egkey, block, stream):
         stream.seek(i)
         return i
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
 
 class NativeTablebases(object):
     """Provides access to Gaviota tablebases via the shared library libgtb."""

diff --git a/test.py b/test.py
@@ -2379,6 +2379,14 @@ def test_wdl(self):
         board = chess.Board("8/8/1p2K3/8/8/3k4/8/8 b - - 0 1")
         self.assertEqual(self.tablebases.probe_wdl(board), 1)
 
+    def test_context_manager(self):
+        self.assertTrue(self.tablebases.available_tables)
+
+        with self.tablebases:
+            pass
+
+        self.assertFalse(self.tablebases.available_tables)
+
 
 if __name__ == "__main__":
     if "-v" in sys.argv or "--verbose" in sys.argv: