Add SafeCloseable context manager wrapper (#5251)

* Add SafeCloseable wrapper * Add SafeCloseable test * Extend JObjectWrapper * Polish * Polish * Apply suggestions from code review Per @chipkent. Co-authored-by: Chip Kent <[email protected]> * Document assumption * Polish pydoc Co-authored-by: Chip Kent <[email protected]> * Change SafeCloseable to AutoCloseable --------- Co-authored-by: Chip Kent <[email protected]>
deephaven · Mar 16, 2024 · d9d4189 · d9d4189
1 parent f1a1373
commit d9d4189
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/py/server/deephaven/jcompat.py b/py/server/deephaven/jcompat.py
@@ -12,7 +12,7 @@
 import pandas as pd
 
 from deephaven import dtypes, DHError
-from deephaven._wrapper import unwrap, wrap_j_object
+from deephaven._wrapper import unwrap, wrap_j_object, JObjectWrapper
 from deephaven.dtypes import DType, _PRIMITIVE_DTYPE_NULL_MAP, _J_ARRAY_NP_TYPE_MAP
 
 _NULL_BOOLEAN_AS_BYTE = jpy.get_type("io.deephaven.util.BooleanUtils").NULL_BOOLEAN_AS_BYTE
@@ -304,3 +304,33 @@ def _j_array_to_series(dtype: DType, j_array: jpy.JType, conv_null: bool) -> pd.
         s = pd.Series(data=np_array, copy=False)
 
     return s
+
+
+class AutoCloseable(JObjectWrapper):
+    """A context manager wrapper to allow Java AutoCloseable to be used in with statements.
+    
+    When constructing a new instance, the Java AutoCloseable must not be closed."""
+
+    j_object_type = jpy.get_type("java.lang.AutoCloseable")
+
+    def __init__(self, j_auto_closeable):
+        self._j_auto_closeable = j_auto_closeable
+        self.closed = False
+
+    def __enter__(self):
+        return self
+
+    def close(self):
+        if not self.closed:
+            self.closed = True
+            self._j_auto_closeable.close()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def __del__(self):
+        self.close()
+
+    @property
+    def j_object(self) -> jpy.JType:
+        return self._j_auto_closeable
diff --git a/py/server/tests/test_jcompat.py b/py/server/tests/test_jcompat.py
@@ -5,11 +5,12 @@
 import unittest
 
 from deephaven import dtypes
-from deephaven.jcompat import j_function, j_lambda
+from deephaven.jcompat import j_function, j_lambda, AutoCloseable
 from tests.testbase import BaseTestCase
 
 import jpy
 
+_JSharedContext = jpy.get_type("io.deephaven.engine.table.SharedContext")
 
 class JCompatTestCase(BaseTestCase):
     def test_j_function(self):
@@ -29,6 +30,12 @@ def int_to_str(v: int) -> str:
         r = j_func.apply(10)
         self.assertEqual(r, "10")
 
+    def test_auto_closeable(self):
+        auto_closeable = AutoCloseable(_JSharedContext.makeSharedContext())
+        with auto_closeable:
+            self.assertEqual(auto_closeable.closed, False)
+        self.assertEqual(auto_closeable.closed, True)
+
 
 if __name__ == "__main__":
     unittest.main()