Create dh.func_in_dtypes from parsing the spec

Author: honno

array_api_tests/dtype_helpers.py

@@ -1,9 +1,12 @@
+import re
 from functools import lru_cache
-from typing import NamedTuple, Tuple, Union
+from inspect import signature
+from typing import Dict, NamedTuple, Tuple, Union
 from warnings import warn
 
 from . import _array_module as xp
 from ._array_module import _UndefinedStub
+from .stubs import name_to_func
 from .typing import DataType, ScalarType
 
 __all__ = [
@@ -99,8 +102,8 @@ class MinMax(NamedTuple):
     xp.uint16: MinMax(0, +65_535),
     xp.uint32: MinMax(0, +4_294_967_295),
     xp.uint64: MinMax(0, +18_446_744_073_709_551_615),
-    xp.float32: MinMax(-3.4028234663852886e+38, 3.4028234663852886e+38),
-    xp.float64: MinMax(-1.7976931348623157e+308, 1.7976931348623157e+308),
+    xp.float32: MinMax(-3.4028234663852886e38, 3.4028234663852886e38),
+    xp.float64: MinMax(-1.7976931348623157e308, 1.7976931348623157e308),
 }
 
 dtype_nbits = {
@@ -196,67 +199,28 @@ def result_type(*dtypes: DataType):
     return result
 
 
-func_in_dtypes = {
-    # elementwise
-    "abs": numeric_dtypes,
-    "acos": float_dtypes,
-    "acosh": float_dtypes,
-    "add": numeric_dtypes,
-    "asin": float_dtypes,
-    "asinh": float_dtypes,
-    "atan": float_dtypes,
-    "atan2": float_dtypes,
-    "atanh": float_dtypes,
-    "bitwise_and": bool_and_all_int_dtypes,
-    "bitwise_invert": bool_and_all_int_dtypes,
-    "bitwise_left_shift": all_int_dtypes,
-    "bitwise_or": bool_and_all_int_dtypes,
-    "bitwise_right_shift": all_int_dtypes,
-    "bitwise_xor": bool_and_all_int_dtypes,
-    "ceil": numeric_dtypes,
-    "cos": float_dtypes,
-    "cosh": float_dtypes,
-    "divide": float_dtypes,
-    "equal": all_dtypes,
-    "exp": float_dtypes,
-    "expm1": float_dtypes,
-    "floor": numeric_dtypes,
-    "floor_divide": numeric_dtypes,
-    "greater": numeric_dtypes,
-    "greater_equal": numeric_dtypes,
-    "isfinite": numeric_dtypes,
-    "isinf": numeric_dtypes,
-    "isnan": numeric_dtypes,
-    "less": numeric_dtypes,
-    "less_equal": numeric_dtypes,
-    "log": float_dtypes,
-    "logaddexp": float_dtypes,
-    "log10": float_dtypes,
-    "log1p": float_dtypes,
-    "log2": float_dtypes,
-    "logical_and": (xp.bool,),
-    "logical_not": (xp.bool,),
-    "logical_or": (xp.bool,),
-    "logical_xor": (xp.bool,),
-    "multiply": numeric_dtypes,
-    "negative": numeric_dtypes,
-    "not_equal": all_dtypes,
-    "positive": numeric_dtypes,
-    "pow": numeric_dtypes,
-    "remainder": numeric_dtypes,
-    "round": numeric_dtypes,
-    "sign": numeric_dtypes,
-    "sin": float_dtypes,
-    "sinh": float_dtypes,
-    "sqrt": float_dtypes,
-    "square": numeric_dtypes,
-    "subtract": numeric_dtypes,
-    "tan": float_dtypes,
-    "tanh": float_dtypes,
-    "trunc": numeric_dtypes,
-    # searching
-    "where": all_dtypes,
+r_in_dtypes = re.compile("x1?: array\n.+Should have an? (.+) data type.")
+r_int_note = re.compile(
+    "If one or both of the input arrays have integer data types, "
+    "the result is implementation-dependent"
+)
+category_to_dtypes = {
+    "boolean": (xp.bool,),
+    "integer": all_int_dtypes,
+    "floating-point": float_dtypes,
+    "numeric": numeric_dtypes,
+    "integer or boolean": bool_and_all_int_dtypes,
 }
+func_in_dtypes: Dict[str, Tuple[DataType, ...]] = {}
+for name, func in name_to_func.items():
+    if m := r_in_dtypes.search(func.__doc__):
+        dtype_category = m.group(1)
+        if dtype_category == "numeric" and r_int_note.search(func.__doc__):
+            dtype_category = "floating-point"
+        dtypes = category_to_dtypes[dtype_category]
+        func_in_dtypes[name] = dtypes
+    elif any("x" in name for name in signature(func).parameters.keys()):
+        func_in_dtypes[name] = all_dtypes
 
 
 func_returns_bool = {