Fix input table jpy restriction

py/server/deephaven/_wrapper.py

@@ -13,7 +13,7 @@
 import pkgutil
 import sys
 from abc import ABC, abstractmethod
-from typing import Set, Union, Optional, Any
+from typing import Set, Union, Optional, Any, List
 
 import jpy
 
@@ -102,21 +102,16 @@ def _is_direct_initialisable(cls) -> bool:
     return False
 
 
-def _lookup_wrapped_class(j_obj: jpy.JType) -> Optional[type]:
-    """ Returns the wrapper class for the specified Java object. """
+def _lookup_wrapped_class(j_obj: jpy.JType) -> List[JObjectWrapper]:
+    """ Returns the wrapper classes for the specified Java object. """
     # load every module in the deephaven package so that all the wrapper classes are loaded and available to wrap
     # the Java objects returned by calling resolve()
     global _has_all_wrappers_imported
     if not _has_all_wrappers_imported:
         _recursive_import(__package__.partition(".")[0])
         _has_all_wrappers_imported = True
 
-    for wc in _di_wrapper_classes:
-        j_clz = wc.j_object_type
-        if j_clz.jclass.isInstance(j_obj):
-            return wc
-
-    return None
+    return [wc for wc in _di_wrapper_classes if wc.j_object_type.jclass.isInstance(j_obj)]
 
 
 def javaify(obj: Any) -> Optional[jpy.JType]:
@@ -162,15 +157,43 @@ def pythonify(j_obj: Any) -> Optional[Any]:
     return wrap_j_object(j_obj)
 
 
-def wrap_j_object(j_obj: jpy.JType) -> Union[JObjectWrapper, jpy.JType]:
-    """ Wraps the specified Java object as an instance of a custom wrapper class if one is available, otherwise returns
-    the raw Java object. """
+def _wrap_with_subclass(j_obj: jpy.JType, cls: type) -> Optional[JObjectWrapper]:
+    """ Returns a wrapper instance for the specified Java object by trying the entire subclasses' hierarchy. The
+    function employs a Depth First Search strategy to try the most specific subclass first. If no matching wrapper class is found,
+    returns None.
+
+    The premises for this function are as follows:
+    - The subclasses all share the same class attribute `j_object_type` (guaranteed by subclassing JObjectWrapper)
+    - The subclasses are all direct initialisable (guaranteed by subclassing JObjectWrapper)
+    - The subclasses are all distinct from each other and check for their uniqueness in the initializer (__init__), e.g.
+      InputTable checks for the presence of the INPUT_TABLE_ATTRIBUTE attribute on the Java object.
+    """
+    for subclass in cls.__subclasses__():
+        try:
+            if (wrapper := _wrap_with_subclass(j_obj, subclass)) is not None:
+                return wrapper
+            return subclass(j_obj)
+        except:
+            continue
+    return None
+
+
+def wrap_j_object(j_obj: jpy.JType) -> Optional[Union[JObjectWrapper, jpy.JType]]:
+    """ Wraps the specified Java object as an instance of the most specific custom wrapper class if one is available,
+    otherwise returns the raw Java object. """
     if j_obj is None:
         return None
 
-    wc = _lookup_wrapped_class(j_obj)
+    wcs = _lookup_wrapped_class(j_obj)
+    for wc in wcs:
+        try:
+            if (wrapper:= _wrap_with_subclass(j_obj, wc)) is not None:
+                return wrapper
+            return wc(j_obj)
+        except:
+            continue
 
-    return wc(j_obj) if wc else j_obj
+    return j_obj
 
 
 def unwrap(obj: Any) -> Union[jpy.JType, Any]:

py/server/deephaven/table_factory.py

@@ -23,7 +23,6 @@
 _JTableFactory = jpy.get_type("io.deephaven.engine.table.TableFactory")
 _JTableTools = jpy.get_type("io.deephaven.engine.util.TableTools")
 _JDynamicTableWriter = jpy.get_type("io.deephaven.engine.table.impl.util.DynamicTableWriter")
-_JBaseArrayBackedInputTable = jpy.get_type("io.deephaven.engine.table.impl.util.BaseArrayBackedInputTable")
 _JAppendOnlyArrayBackedInputTable = jpy.get_type(
     "io.deephaven.engine.table.impl.util.AppendOnlyArrayBackedInputTable")
 _JKeyedArrayBackedInputTable = jpy.get_type("io.deephaven.engine.table.impl.util.KeyedArrayBackedInputTable")
@@ -235,7 +234,6 @@ class InputTable(Table):
 
     Users should always create InputTables through factory methods rather than directly from the constructor.
     """
-    j_object_type = _JBaseArrayBackedInputTable
 
     def __init__(self, j_table: jpy.JType):
         super().__init__(j_table)

py/server/tests/test_table_factory.py

@@ -9,18 +9,19 @@
 import numpy as np
 
 from deephaven import DHError, read_csv, time_table, empty_table, merge, merge_sorted, dtypes, new_table, \
-    input_table, time
+    input_table, time, _wrapper
 from deephaven.column import byte_col, char_col, short_col, bool_col, int_col, long_col, float_col, double_col, \
     string_col, datetime_col, pyobj_col, jobj_col
 from deephaven.constants import NULL_DOUBLE, NULL_FLOAT, NULL_LONG, NULL_INT, NULL_SHORT, NULL_BYTE
-from deephaven.table_factory import DynamicTableWriter, ring_table
+from deephaven.table_factory import DynamicTableWriter, InputTable, ring_table
 from tests.testbase import BaseTestCase
 from deephaven.table import Table
 from deephaven.stream import blink_to_append_only, stream_to_append_only
 
 JArrayList = jpy.get_type("java.util.ArrayList")
 _JBlinkTableTools = jpy.get_type("io.deephaven.engine.table.impl.BlinkTableTools")
 _JDateTimeUtils = jpy.get_type("io.deephaven.time.DateTimeUtils")
+_JTable = jpy.get_type("io.deephaven.engine.table.Table")
 
 
 @dataclass
@@ -372,6 +373,14 @@ def test_input_table(self):
             keyed_input_table.delete(t.select(["String", "Double"]))
             self.assertEqual(keyed_input_table.size, 0)
 
+        with self.subTest("custom input table creation"):
+            place_holder_input_table = empty_table(1).update_view(["Key=`A`", "Value=10"]).with_attributes({_JTable.INPUT_TABLE_ATTRIBUTE: "Placeholder IT"}).j_table
+            # Confirming no error.
+            it = InputTable(place_holder_input_table)
+
+            self.assertTrue(isinstance(_wrapper.wrap_j_object(place_holder_input_table), InputTable))
+
+
     def test_ring_table(self):
         cols = [
             bool_col(name="Boolean", data=[True, False]),
@@ -450,5 +459,18 @@ def test_input_table_empty_data(self):
         it.delete(t)
         self.assertEqual(it.size, 0)
 
+    def test_j_input_wrapping(self):
+        cols = [
+            bool_col(name="Boolean", data=[True, False]),
+            string_col(name="String", data=["foo", "bar"]),
+        ]
+        t = new_table(cols=cols)
+        col_defs = {c.name: c.data_type for c in t.columns}
+        append_only_input_table = input_table(col_defs=col_defs)
+
+        t = _wrapper.wrap_j_object(append_only_input_table.j_table)
+        self.assertTrue(isinstance(t, InputTable))
+
+
 if __name__ == '__main__':
     unittest.main()