The return of ND arrays!

gselzer · ctrueden · commit 87232a8ab51c · 2022-11-18T08:50:06.000-06:00
They're slower than primitive arrays, but better to have slow arrays
than no arrays?
diff --git a/src/scyjava/_java.py b/src/scyjava/_java.py
@@ -523,11 +523,30 @@ def jarray(kind, lengths: Sequence):
     if mode == Mode.JEP:
         import jep  # noqa: F401
 
-        # TODO: Support n-d arrays
-        if len(lengths) > 1:
-            raise RuntimeError("jep cannot support 2+ dimensional arrays!")
-        # instantiate the n-dimensional array
-        arr = jep.jarray(lengths[0], arraytype)
+        if len(lengths) == 1:
+            # Fast case: 1-d array (we can use primitives)
+            arr = jep.jarray(lengths[0], arraytype)
+        else:
+            # Slow case: n-d array (we cannot use primitives)
+            # See https://github.com/ninia/jep/issues/439
+            kinds = {
+                "b": jimport("java.lang.Byte"),
+                "c": jimport("java.lang.Character"),
+                "d": jimport("java.lang.Double"),
+                "f": jimport("java.lang.Float"),
+                "i": jimport("java.lang.Integer"),
+                "j": jimport("java.lang.Long"),
+                "s": jimport("java.lang.Short"),
+                "z": jimport("java.lang.Boolean"),
+            }
+            if arraytype in kinds:
+                arraytype = kinds[arraytype]
+                kind = arraytype
+            # build up the array type
+            for _ in range(len(lengths) - 1):
+                arraytype = jep.jarray(0, arraytype)
+            # instantiate the n-dimensional array
+            arr = jep.jarray(lengths[0], arraytype)
 
     elif mode == Mode.JPYPE:
         start_jvm()
diff --git a/tests/test_arrays.py b/tests/test_arrays.py
@@ -1,5 +1,4 @@
 import numpy as np
-import pytest
 
 from scyjava import is_jarray, jarray, to_python
 from scyjava.config import Mode, mode
@@ -47,9 +46,6 @@ def assert_array_conversion_works(jarr, expected):
         assert_array_conversion_works(jints, deltas)
 
     def test_jarray2d_to_python(self):
-        if mode is Mode.JEP:
-            pytest.skip("Jep doesn't support 2-d arrays")
-
         nums = [
             [1.2, 3.4, 5.6],
             [7.8, 9.1, 2.3],
@@ -69,7 +65,9 @@ def test_jarray2d_to_python(self):
         pdoubles = to_python(jdoubles)
 
         if mode == Mode.JEP:
-            raise RuntimeError("Not supported")
+            assert isinstance(pdoubles, list)
+            assert all(isinstance(v, list) for v in pdoubles)
+            assert len(nums) == len(pdoubles)
 
         elif mode == Mode.JPYPE:
             assert isinstance(pdoubles, np.ndarray)
@@ -81,9 +79,6 @@ def test_jarray2d_to_python(self):
                 assert nums[i][j] == pdoubles[i][j]
 
     def test_jarray2d_to_python_updates(self):
-        if mode is Mode.JEP:
-            pytest.skip("Jep doesn't support 2-d arrays")
-
         nums_init = [
             [1.2, 3.4, 5.6],
             [7.8, 9.1, 2.3],
@@ -105,9 +100,15 @@ def test_jarray2d_to_python_updates(self):
 
         # assert narr initial state
         pdoubles = to_python(jdoubles)
-        assert isinstance(pdoubles, np.ndarray)
-        assert np.float64 == pdoubles.dtype
-        assert (5, 3) == pdoubles.shape
+        if mode == Mode.JEP:
+            assert isinstance(pdoubles, list)
+            assert isinstance(pdoubles[0][0], float)
+            assert len(pdoubles) == 5
+            assert len(pdoubles[0]) == 3
+        elif mode == Mode.JPYPE:
+            assert isinstance(pdoubles, np.ndarray)
+            assert np.float64 == pdoubles.dtype
+            assert (5, 3) == pdoubles.shape
         for i in range(len(nums_init)):
             for j in range(len(nums_init[i])):
                 assert nums_init[i][j] == pdoubles[i][j]
@@ -119,9 +120,15 @@ def test_jarray2d_to_python_updates(self):
 
         # assert narr delta state
         pdoubles = to_python(jdoubles)
-        assert isinstance(pdoubles, np.ndarray)
-        assert np.float64 == pdoubles.dtype
-        assert (5, 3) == pdoubles.shape
+        if mode == Mode.JEP:
+            assert isinstance(pdoubles, list)
+            assert isinstance(pdoubles[0][0], float)
+            assert len(pdoubles) == 5
+            assert len(pdoubles[0]) == 3
+        elif mode == Mode.JPYPE:
+            assert isinstance(pdoubles, np.ndarray)
+            assert np.float64 == pdoubles.dtype
+            assert (5, 3) == pdoubles.shape
         for i in range(len(nums_delta)):
             for j in range(len(nums_delta[i])):
                 assert nums_delta[i][j] == pdoubles[i][j]
