More pathological test cases

py/server/tests/test_udf_numpy_args.py

@@ -1,14 +1,15 @@
 #
 #     Copyright (c) 2016-2023 Deephaven Data Labs and Patent Pending
 #
-
+import typing
+from typing import Optional, Union, Any
 import unittest
 
-from typing import Optional
 
 import numpy as np
+import numpy.typing as npt
 
-from deephaven import empty_table, DHError
+from deephaven import empty_table, DHError, dtypes
 from deephaven.dtypes import double_array, int32_array, long_array, int16_array, char_array, int8_array, \
     float32_array
 from tests.testbase import BaseTestCase
@@ -179,6 +180,70 @@ def test_udf(col: Optional[{np_type}]) -> bool:
                     res = tbl.update("Col3 = test_udf(Col2)")
                     self.assertEqual(4, res.to_string().count("true"))
 
+    def test_weird_cases(self):
+        def f(p1: Union[np.ndarray[typing.Any], None]) -> bool:
+            return bool(p1)
+        with self.assertRaises(DHError) as cm:
+            t = empty_table(10).update(["X1 = f(i)"])
+
+        def f1(p1: Union[np.int16, np.int32]) -> bool:
+            return bool(p1)
+        with self.assertRaises(DHError) as cm:
+            t = empty_table(10).update(["X1 = f1(i)"])
+
+        def f11(p1: Union[float, np.float32]) -> bool:
+            return bool(p1)
+        with self.assertRaises(DHError) as cm:
+            t = empty_table(10).update(["X1 = f11(i)"])
+
+        def f2(p1: Union[np.int16, np.float64]) -> Union[Optional[bool]]:
+            return bool(p1)
+        t = empty_table(10).update(["X1 = f2(i)"])
+        self.assertEqual(t.columns[0].data_type, dtypes.bool_)
+        self.assertEqual(9, t.to_string().count("true"))
+
+        def f21(p1: Union[np.int16, np.float64]) -> Union[Optional[bool], int]:
+            return bool(p1)
+        with self.assertRaises(DHError) as cm:
+            t = empty_table(10).update(["X1 = f21(i)"])
+
+        def f3(p1: Union[np.int16, np.float64], p2=None) -> bool:
+            return bool(p1)
+        t = empty_table(10).update(["X1 = f3(i)"])
+        self.assertEqual(t.columns[0].data_type, dtypes.bool_)
+
+        def f4(p1: Union[np.int16, np.float64], p2=None) -> bool:
+            return bool(p1)
+        t = empty_table(10).update(["X1 = f4((double)i)"])
+        self.assertEqual(t.columns[0].data_type, dtypes.bool_)
+        with self.assertRaises(DHError) as cm:
+            t = empty_table(10).update(["X1 = f4(now())"])
+
+        def f41(p1: Union[np.int16, np.float64, Union[Any]], p2=None) -> bool:
+            return bool(p1)
+        t = empty_table(10).update(["X1 = f41(now())"])
+        self.assertEqual(t.columns[0].data_type, dtypes.bool_)
+
+        def f42(p1: Union[np.int16, np.float64, np.datetime64], p2=None) -> bool:
+            return p1.dtype.char == "M"
+        t = empty_table(10).update(["X1 = f42(now())"])
+        self.assertEqual(t.columns[0].data_type, dtypes.bool_)
+        self.assertEqual(10, t.to_string().count("true"))
+
+        def f5(col1, col2: np.ndarray[np.int32]) -> bool:
+            return np.nanmean(col2) == np.mean(col2)
+        t = empty_table(10).update(["X = i % 3", "Y = i"]).group_by("X")
+        t = t.update(["X1 = f5(X, Y)"])
+        with self.assertRaises(DHError) as cm:
+            t = t.update(["X1 = f5(X, null)"])
+
+        def f51(col1, col2: Optional[np.ndarray[np.int32]]) -> bool:
+            return np.nanmean(col2) == np.mean(col2)
+        t = empty_table(10).update(["X = i % 3", "Y = i"]).group_by("X")
+        t = t.update(["X1 = f51(X, Y)"])
+        with self.assertRaises(DHError) as cm:
+            t = t.update(["X1 = f51(X, null)"])
+
 
 if __name__ == "__main__":
     unittest.main()