From d9d4189384df7651be2961f0ba30eccb154e0cb3 Mon Sep 17 00:00:00 2001
From: arman-ddl <122062464+arman-ddl@users.noreply.github.com>
Date: Fri, 15 Mar 2024 22:34:14 -0700
Subject: [PATCH] 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 <5250374+chipkent@users.noreply.github.com>

* Document assumption

* Polish pydoc

Co-authored-by: Chip Kent <5250374+chipkent@users.noreply.github.com>

* Change SafeCloseable to AutoCloseable

---------

Co-authored-by: Chip Kent <5250374+chipkent@users.noreply.github.com>
---
 py/server/deephaven/jcompat.py  | 32 +++++++++++++++++++++++++++++++-
 py/server/tests/test_jcompat.py |  9 ++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/py/server/deephaven/jcompat.py b/py/server/deephaven/jcompat.py
index a4fb34c2495..46cb3d79a3a 100644
--- 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
index bcea58d112c..40241b73da2 100644
--- 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()