test_signatures.py

"""
Tests for function/method signatures compliance

We're not interested in being 100% strict - instead we focus on areas which
could affect interop, e.g. with

 def add(x1, x2, /):
 ...

x1 and x2 don't need to be pos-only for the purposes of interoperability, but with

 def squeeze(x, /, axis):
 ...

axis has to be pos-or-keyword to support both styles

 >>> squeeze(x, 0)
 ...
 >>> squeeze(x, axis=0)
 ...

"""
fromcollectionsimportdefaultdict
fromcopyimportcopy
frominspectimportParameter, Signature, signature
fromtypesimportFunctionType
fromtypingimportAny, Callable, Dict, Literal, get_args
fromwarningsimportwarn

importpytest

from . importdtype_helpersasdh
from . importxp
from .stubsimport (array_methods, category_to_funcs, extension_to_funcs,
name_to_func, info_funcs)

ParameterKind=Literal[
Parameter.POSITIONAL_ONLY,
Parameter.VAR_POSITIONAL,
Parameter.POSITIONAL_OR_KEYWORD,
Parameter.KEYWORD_ONLY,
Parameter.VAR_KEYWORD,
]
ALL_KINDS=get_args(ParameterKind)
VAR_KINDS= (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD)
kind_to_str: Dict[ParameterKind, str] = {
Parameter.POSITIONAL_OR_KEYWORD: "pos or kw argument",
Parameter.POSITIONAL_ONLY: "pos-only argument",
Parameter.KEYWORD_ONLY: "keyword-only argument",
Parameter.VAR_POSITIONAL: "star-args (i.e. *args) argument",
Parameter.VAR_KEYWORD: "star-kwargs (i.e. **kwargs) argument",
}


def_test_inspectable_func(sig: Signature, stub_sig: Signature):
params=list(sig.parameters.values())
stub_params=list(stub_sig.parameters.values())

non_kwonly_stub_params= [
pforpinstub_paramsifp.kind!=Parameter.KEYWORD_ONLY
 ]
# sanity check
assertnon_kwonly_stub_params==stub_params[: len(non_kwonly_stub_params)]
# We're not interested if the array module has additional arguments, so we
# only iterate through the arguments listed in the spec.
fori, stub_paraminenumerate(non_kwonly_stub_params):
assert (
len(params) >=i+1
 ), f"Argument '{stub_param.name}' missing from signature"
param=params[i]

# We're not interested in the name if it isn't actually used
ifstub_param.kindnotin [Parameter.POSITIONAL_ONLY, *VAR_KINDS]:
assert (
param.name==stub_param.name
 ), f"Expected argument '{param.name}' to be named '{stub_param.name}'"

ifstub_param.kindin [Parameter.POSITIONAL_OR_KEYWORD, *VAR_KINDS]:
f_stub_kind=kind_to_str[stub_param.kind]
assertparam.kind==stub_param.kind, (
f"{param.name} is a {kind_to_str[param.kind]}, "
f"but should be a {f_stub_kind}"
 )

kwonly_stub_params=stub_params[len(non_kwonly_stub_params) :]
forstub_paraminkwonly_stub_params:
assert (
stub_param.nameinsig.parameters.keys()
 ), f"Argument '{stub_param.name}' missing from signature"
param=next(pforpinparamsifp.name==stub_param.name)
f_stub_kind=kind_to_str[stub_param.kind]
assertparam.kindin [stub_param.kind, Parameter.POSITIONAL_OR_KEYWORD,], (
f"{param.name} is a {kind_to_str[param.kind]}, "
f"but should be a {f_stub_kind} "
f"(or at least a {kind_to_str[ParameterKind.POSITIONAL_OR_KEYWORD]})"
 )


defmake_pretty_func(func_name: str, *args: Any, **kwargs: Any) ->str:
f_sig=f"{func_name}("
f_sig+=", ".join(str(a) forainargs)
iflen(kwargs) !=0:
iflen(args) !=0:
f_sig+=", "
f_sig+=", ".join(f"{k}={v}"fork, vinkwargs.items())
f_sig+=")"
returnf_sig


# We test uninspectable signatures by passing valid, manually-defined arguments
# to the signature's function/method.
#
# Arguments which require use of the array module are specified as string
# expressions to be eval()'d on runtime. This is as opposed to just using the
# array module whilst setting up the tests, which is prone to halt the entire
# test suite if an array module doesn't support a given expression.
func_to_specified_args=defaultdict(
dict,
 {
"permute_dims": {"axes": 0},
"reshape": {"shape": (1, 5)},
"broadcast_to": {"shape": (1, 5)},
"asarray": {"obj": [0, 1, 2, 3, 4]},
"full_like": {"fill_value": 42},
"matrix_power": {"n": 2},
 },
)
func_to_specified_arg_exprs=defaultdict(
dict,
 {
"stack": {"arrays": "[xp.ones((5,)), xp.ones((5,))]"},
"iinfo": {"type": "xp.int64"},
"finfo": {"type": "xp.float64"},
"cholesky": {"x": "xp.asarray([[1, 0], [0, 1]], dtype=xp.float64)"},
"inv": {"x": "xp.asarray([[1, 2], [3, 4]], dtype=xp.float64)"},
"solve": {
a: "xp.asarray([[1, 2], [3, 4]], dtype=xp.float64)"forain ["x1", "x2"]
 },
"outer": {"x1": "xp.ones((5,))", "x2": "xp.ones((5,))"},
 },
)
# We default most array arguments heuristically. As functions/methods work only
# with arrays of certain dtypes and shapes, we specify only supported arrays
# respective to the function.
casty_names= ["__bool__", "__int__", "__float__", "__complex__", "__index__"]
matrixy_names= [
f.__name__
forfincategory_to_funcs["linear_algebra"] +extension_to_funcs["linalg"]
]
matrixy_names+= ["__matmul__", "triu", "tril"]
forfunc_name, funcinname_to_func.items():
stub_sig=signature(func)
array_argnames=set(stub_sig.parameters.keys()) & {"x", "x1", "x2", "other"}
iffuncinarray_methods:
array_argnames.add("self")
array_argnames-=set(func_to_specified_arg_exprs[func_name].keys())
iflen(array_argnames) >0:
in_dtypes=dh.func_in_dtypes[func_name]
fordtype_namein ["float64", "bool", "int64", "complex128"]:
# We try float64 first because uninspectable numerical functions
# tend to support float inputs first-and-foremost (i.e. PyTorch)
try:
dtype=getattr(xp, dtype_name)
exceptAttributeError:
pass
else:
ifdtypeinin_dtypes:
iffunc_nameincasty_names:
shape= ()
eliffunc_nameinmatrixy_names:
shape= (3, 3)
else:
shape= (5,)
fallback_array_expr=f"xp.ones({shape}, dtype=xp.{dtype_name})"
break
else:
warn(
f"{dh.func_in_dtypes['{func_name}']}={in_dtypes} seemingly does "
"not contain any assumed dtypes, so skipping specifying fallback array."
 )
continue
forargnameinarray_argnames:
func_to_specified_arg_exprs[func_name][argname] =fallback_array_expr


def_test_uninspectable_func(func_name: str, func: Callable, stub_sig: Signature):
params=list(stub_sig.parameters.values())

iflen(params) ==0:
func()
return

uninspectable_msg= (
f"Note {func_name}() is not inspectable so arguments are passed "
"manually to test the signature."
 )

argname_to_arg=copy(func_to_specified_args[func_name])
argname_to_expr=func_to_specified_arg_exprs[func_name]
forargname, exprinargname_to_expr.items():
assertargnamenotinargname_to_arg.keys() # sanity check
try:
argname_to_arg[argname] =eval(expr, {"xp": xp})
exceptExceptionase:
pytest.skip(
f"Exception occured when evaluating {argname}={expr}: {e}\n"
f"{uninspectable_msg}"
 )

posargs= []
posorkw_args= {}
kwargs= {}
no_arg_msg= (
"We have no argument specified for '{}'. Please ensure you're using "
"the latest version of array-api-tests, then open an issue if one "
f"doesn't already exist. {uninspectable_msg}"
 )
forparaminparams:
ifparam.kind==Parameter.POSITIONAL_ONLY:
try:
posargs.append(argname_to_arg[param.name])
exceptKeyError:
pytest.skip(no_arg_msg.format(param.name))
elifparam.kind==Parameter.POSITIONAL_OR_KEYWORD:
ifparam.default==Parameter.empty:
try:
posorkw_args[param.name] =argname_to_arg[param.name]
exceptKeyError:
pytest.skip(no_arg_msg.format(param.name))
else:
assertargname_to_arg[param.name]
posorkw_args[param.name] =param.default
elifparam.kind==Parameter.KEYWORD_ONLY:
assertparam.default!=Parameter.empty# sanity check
kwargs[param.name] =param.default
else:
assertparam.kindinVAR_KINDS# sanity check
pytest.skip(no_arg_msg.format(param.name))
iflen(posorkw_args) ==0:
func(*posargs, **kwargs)
else:
posorkw_name_to_arg_pairs=list(posorkw_args.items())
foriinrange(len(posorkw_name_to_arg_pairs), -1, -1):
extra_posargs= [argfor_, arginposorkw_name_to_arg_pairs[:i]]
extra_kwargs=dict(posorkw_name_to_arg_pairs[i:])
func(*posargs, *extra_posargs, **kwargs, **extra_kwargs)


def_test_func_signature(func: Callable, stub: FunctionType, is_method=False):
stub_sig=signature(stub)
# If testing against array, ignore 'self' arg in stub as it won't be present
# in func (which should be a method).
ifis_method:
stub_params=list(stub_sig.parameters.values())
ifstub_params[0].name=="self":
delstub_params[0]
stub_sig=Signature(
parameters=stub_params, return_annotation=stub_sig.return_annotation
 )

try:
sig=signature(func)
exceptValueError:
try:
_test_uninspectable_func(stub.__name__, func, stub_sig)
exceptExceptionase:
raiseefromNone# suppress parent exception for cleaner pytest output
else:
_test_inspectable_func(sig, stub_sig)


@pytest.mark.parametrize(
"stub",
 [sforstubsincategory_to_funcs.values() forsinstubs],
ids=lambdaf: f.__name__,
)
deftest_func_signature(stub: FunctionType):
asserthasattr(xp, stub.__name__), f"{stub.__name__} not found in array module"
func=getattr(xp, stub.__name__)
_test_func_signature(func, stub)


extension_and_stub_params= []
forext, stubsinextension_to_funcs.items():
forstubinstubs:
p=pytest.param(
ext, stub, id=f"{ext}.{stub.__name__}", marks=pytest.mark.xp_extension(ext)
 )
extension_and_stub_params.append(p)


@pytest.mark.parametrize("extension, stub", extension_and_stub_params)
deftest_extension_func_signature(extension: str, stub: FunctionType):
mod=getattr(xp, extension)
asserthasattr(
mod, stub.__name__
 ), f"{stub.__name__} not found in {extension} extension"
func=getattr(mod, stub.__name__)
_test_func_signature(func, stub)


@pytest.mark.parametrize("stub", array_methods, ids=lambdaf: f.__name__)
deftest_array_method_signature(stub: FunctionType):
x_expr=func_to_specified_arg_exprs[stub.__name__]["self"]
try:
x=eval(x_expr, {"xp": xp})
exceptExceptionase:
pytest.skip(f"Exception occured when evaluating x={x_expr}: {e}")
asserthasattr(x, stub.__name__), f"{stub.__name__} not found in array object {x!r}"
method=getattr(x, stub.__name__)
_test_func_signature(method, stub, is_method=True)

ifinfo_funcs: # pytest fails collecting if info_funcs is empty
@pytest.mark.min_version("2023.12")
@pytest.mark.parametrize("stub", info_funcs, ids=lambdaf: f.__name__)
deftest_info_func_signature(stub: FunctionType):
try:
info_namespace=xp.__array_namespace_info__()
exceptExceptionase:
raiseAssertionError(f"Could not get info namespace from xp.__array_namespace_info__(): {e}")

func=getattr(info_namespace, stub.__name__)
_test_func_signature(func, stub)