casting_assumptions.py

import math
import sys
from itertools import product
from operator import index

import numpy as np
import pytest
from hypothesis import assume, given, note
from hypothesis import strategies as st
from hypothesis.extra import numpy as nps

assert sys.maxsize == 2**63 - 1  # assume 64-bit Python

builtin_to_dtype = {
    bool: np.bool_,
    int: np.int64,
    float: np.float64,
    complex: np.complex128,
}


@pytest.mark.parametrize("use_np", [False, True])
@pytest.mark.parametrize(
    "builtin, arg_kind", product([bool, int, float, complex], repeat=2)
)
@given(data=st.data())
def test_builtin_casting(builtin, arg_kind, use_np, data):
    if use_np:
        _math = np
        dtype = builtin_to_dtype[arg_kind]
        # We forcibly map examples with np.asarray to ensure we don't get NumPy
        # scalars. This is because they have slightly different behaviour, e.g.
        #
        #     >>> np_array = np.asarray(1+1j, dtype=np.complex128)
        #     >>> int(np_array)
        #     TypeError: can't convert complex to int
        #     >>> np_scalar = np.complex128(1+1j)
        #     >>> int(np_scalar)
        #     1
        #
        arg_strat = nps.from_dtype(np.dtype(dtype)).map(np.asarray)
    else:
        _math = math
        arg_strat = st.from_type(arg_kind)
    arg = data.draw(arg_strat, label="arg")
    f_expr = f"{builtin.__name__}({arg})"
    if builtin == int and arg_kind == float and not _math.isfinite(arg):
        if _math.isinf(arg):
            errclass = OverflowError
        else:
            assert _math.isnan(arg)  # sanity check
            errclass = ValueError
        with pytest.raises(errclass):
            builtin(arg)
        note(f"{f_expr}=<{errclass.__name__}>")
    elif builtin in [int, float] and arg_kind == complex:
        with pytest.raises(TypeError):
            builtin(arg)
        note(f"{f_expr}=<TypeError>")
    else:
        res = builtin(arg)
        note(f"{f_expr}={res}")
        assert type(res) == builtin
        if builtin == bool:
            assert not res if arg == 0 else res
        elif arg_kind == bool:
            assert res == 1 if arg else res == 0
        elif builtin == int:
            assert arg_kind in [int, float]  # sanity check
            if arg_kind == float:
                assume(abs(arg) < 2**63)
            assert res == _math.trunc(arg)
        elif builtin == float:
            assert arg_kind in [int, float]  # sanity check
            if _math.isnan(arg):
                assert _math.isnan(res)
            else:
                if arg_kind == int:
                    assume(abs(arg) <= 2**53)
                assert res == arg
        else:
            assert builtin == complex  # sanity check
            if arg_kind in [int, float]:
                assert res.imag == 0
                if _math.isnan(arg):
                    assert _math.isnan(res.real)
                else:
                    if arg_kind == int:
                        assume(abs(arg) <= 2**53)
                    assert res.real == arg
            else:
                assert arg_kind == complex  # sanity check
                for component, arg_component in [
                    (res.real, arg.real),
                    (res.imag, arg.imag),
                ]:
                    if _math.isnan(arg_component):
                        assert _math.isnan(component)
                    else:
                        assert component == arg_component


@pytest.mark.parametrize("use_np", [False, True])
@pytest.mark.parametrize("arg_kind", [bool, int, float, complex])
@given(data=st.data())
def test_index(arg_kind, use_np, data):
    if use_np:
        _math = np
        dtype = builtin_to_dtype[arg_kind]
        # See comment in test_builtin_casting
        arg_strat = nps.from_dtype(np.dtype(dtype)).map(np.asarray)
    else:
        _math = math
        arg_strat = st.from_type(arg_kind)
    arg = data.draw(arg_strat, label="arg")
    f_expr = f"index({arg})"
    if arg_kind in [float, complex] or (arg_kind == bool and use_np):
        with pytest.raises(TypeError):
            index(arg)
        note(f"{f_expr}=<TypeError>")
    else:
        res = index(arg)
        note(f"{f_expr}={res}")
        assert type(res) == bool if arg_kind == bool else int
        if arg_kind == bool:
            assert res is True if arg else res is False
        else:
            assert arg_kind == int  # sanity check
            assert res == _math.trunc(arg)