diff --git a/py/server/deephaven/_udf.py b/py/server/deephaven/_udf.py index fba76b3472a..3ec5b93f95a 100644 --- a/py/server/deephaven/_udf.py +++ b/py/server/deephaven/_udf.py @@ -240,6 +240,14 @@ def _parse_signature(fn: Callable) -> _ParsedSignature: return p_sig +def _is_from_np_type(param_types:set[type], np_type_char: str) -> bool: + """ Determine if the given numpy type char comes for a numpy type in the given set of parameter type annotations""" + for t in param_types: + if issubclass(t, np.generic) and np.dtype(t).char == np_type_char: + return True + return False + + def _convert_arg(param: _ParsedParamAnnotation, arg: Any) -> Any: """ Convert a single argument to the type specified by the annotation """ if arg is None: @@ -279,13 +287,21 @@ def _convert_arg(param: _ParsedParamAnnotation, arg: Any) -> Any: else: raise DHError(f"Argument {arg} is not compatible with annotation {param.orig_types}") else: - return np.dtype(param.int_char).type(arg) + # return a numpy integer instance only if the annotation is a numpy type + if _is_from_np_type(param.orig_types, param.int_char): + return np.dtype(param.int_char).type(arg) + else: + return arg elif param.floating_char and isinstance(arg, float): if isinstance(arg, float): if arg == dh_null: return np.nan if "N" not in param.encoded_types else None else: - return np.dtype(param.floating_char).type(arg) + # return a numpy floating instance only if the annotation is a numpy type + if _is_from_np_type(param.orig_types, param.floating_char): + return np.dtype(param.floating_char).type(arg) + else: + return arg elif t == "?" and isinstance(arg, bool): return arg elif t == "M": diff --git a/py/server/tests/test_udf_numpy_args.py b/py/server/tests/test_udf_numpy_args.py index ba698a4b21c..78544bf8ec2 100644 --- a/py/server/tests/test_udf_numpy_args.py +++ b/py/server/tests/test_udf_numpy_args.py @@ -392,6 +392,41 @@ def f31(p1: Optional[np.bool_], p2=None) -> bool: t2 = t.update(["X1 = f31(null, Y)"]) self.assertEqual(10, t2.to_string("X1").count("true")) + def test_non_np_typehints(self): + py_types = {"int", "float"} + + for p_type in py_types: + with self.subTest(p_type): + func_str = f""" +def f(x: {p_type}) -> bool: # note typing + return type(x) == {p_type} +""" + exec(func_str, globals()) + t = empty_table(1).update(["X = i", f"Y = f(({p_type})X)"]) + self.assertEqual(1, t.to_string(cols="Y").count("true")) + + + np_int_types = {"np.int8", "np.int16", "np.int32", "np.int64"} + for p_type in np_int_types: + with self.subTest(p_type): + func_str = f""" +def f(x: {p_type}) -> bool: # note typing + return type(x) == {p_type} +""" + exec(func_str, globals()) + t = empty_table(1).update(["X = i", f"Y = f(X)"]) + self.assertEqual(1, t.to_string(cols="Y").count("true")) + + np_floating_types = {"np.float32", "np.float64"} + for p_type in np_floating_types: + with self.subTest(p_type): + func_str = f""" +def f(x: {p_type}) -> bool: # note typing + return type(x) == {p_type} +""" + exec(func_str, globals()) + t = empty_table(1).update(["X = i", f"Y = f((float)X)"]) + self.assertEqual(1, t.to_string(cols="Y").count("true")) if __name__ == "__main__": unittest.main()