def test_array_and_op(): array = zn.Array(zn.var(int), shape=(3, 4)) v = zn.var(int) iter_var, op = get_iter_var_and_op(array, v + 1) assert op.op == _Op_code.add assert isinstance(op.params[1], int) and op.params[1] == 1 assert isinstance(op.params[0], zn.var) and op.params[0] is v
class TestToStr: @pytest.mark.parametrize("name, arg", [("a", (1, 2)), ("b", zn.var(range(1, 5)))]) def test_name(self, name, arg): a = zn.Set(arg) a._name = name assert name == to_str(a)
def test_slice_wrong_start(self, start): array = zn.Array(zn.var(int), shape=(3, 2, 4)) with pytest.raises( ValueError, match=re.escape( f"start({start}) should be smaller then stop(10)")): _ = array[1, start:10]
class TestSize: @pytest.mark.parametrize("dim", (-1, 3)) def test_wrong_dim(self, dim): a = zn.Array([[1, 2], [2, 3], [3, 4]]) with pytest.raises( ValueError, match=f"Array has 0\\.\\.2 dimensions, but {dim} were specified" ): a.size(dim) @pytest.mark.parametrize("array", [ zn.Array([[1, 2], [2, 3], [3, 4]]), zn.Array(zn.var(int), shape=(3, 2)) ]) @pytest.mark.parametrize("dim, expected", ((0, 3), (1, 2))) def test_dim(self, array, dim, expected): class MyModel(zn.Model): def __init__(self, array): self.a = array self.s = self.a.size(dim) self.s2 = self.a.size(dim) * 2 # test arithmetic operations self.s_bigger_2 = self.a.size( dim) > 2 # test compare operations assert self.s.type is int assert self.s2.type is int assert self.s_bigger_2.type is int result = MyModel(array).solve_satisfy() assert result["s"] == expected assert result["s_bigger_2"] == (expected > 2) assert result["s2"] == (expected * 2)
def _(seq: Union[list, tuple]): if seq: v = var(get_type(seq[0])) else: raise ValueError( "empty sequences are not supported as constraint parameter") return v
def __init__(self, a: int, b: int, c: int): super().__init__() self.a = par(a) self.b = par(b) self.c = par(c) self.x = var(range(-100, 101)) self.constraints = [ self.a * self.x**2 + self.b * self.x + self.c == 0 ]
def test_slice_and_func(): def fn(par): return 2 - par array = zn.Array(zn.var(int), shape=(3, 4)) iter_var, op = get_iter_var_and_op(array[2:], fn) assert iter_var.name == "par" assert iter_var.type is int assert op.op == _Op_code.sub assert isinstance(op.params[0], int) and op.params[0] == 2 assert op.params[1].name == "par"
class TestCreate: def test_int_tuple(self): a = zn.Set((1, 2, 3)) assert isinstance(a, SetPar) assert a.value == (1, 2, 3) @pytest.mark.parametrize("collection", [zn.range(1, 4), range(1, 4)]) def test_int_range(self, collection): a = zn.Set(collection) assert isinstance(a, SetPar) assert a.value == collection def test_int_generator(self): a = zn.Set((a + 1 for a in range(4))) assert isinstance(a, SetPar) assert a.value == (1, 2, 3, 4) @pytest.mark.parametrize( "arg", ["aaa", zn.var(float), (a for a in "aaa"), (1.2, 2.5)]) def test_wrong_type(self, arg): with pytest.raises(ValueError, match="Unsupported type for set: <class"): zn.Set(arg)
def _(seq: ArrayMixin): return var(seq.type)
class TestBinary: v = zn.var(int) v._name = "i" p = zn.par(2) p._name = "a" variables = v, p @pytest.mark.parametrize("v", variables) @pytest.mark.parametrize( "op, sign", ([operator.add, "+"], [operator.sub, "-"], [operator.floordiv, "div"], [operator.mul, "*"], [operator.mod, "mod"], [operator.eq, "=="], [operator.ne, "!="], [operator.gt, ">"], [operator.lt, "<"], [operator.ge, ">="], [operator.le, "<="])) def test_simple_ariph_binary_op(self, v, op: Callable, sign: str): result = op(v, 1) assert to_str(result) == f"({v.name} {sign} 1)" assert result.type is int def test_complex_binary_op(self): result = (1 - self.v) * (2 // self.p) >= 4 assert to_str( result) == f"(((1 - {self.v.name}) * (2 div {self.p.name})) >= 4)" assert result.type is int def test_long_binary_op(self): result = 1 + self.v * 2 - self.p + self.v**self.p assert to_str( result ) == f"(((1 + ({self.v.name} * 2)) - {self.p.name}) + pow({self.v.name}, {self.p.name}))" assert result.type is int def test_pow_bracket(self): result = -3 * self.v**(2 % self.p) assert to_str( result) == f"(-3 * pow({self.v.name}, (2 mod {self.p.name})))" assert result.type is int def test_wrong_pow(self): with pytest.raises(ValueError, match="modulo is not supported"): pow(self.p, 2, 10) def test_wrong_div(self): with pytest.raises(ValueError, match="right part of expression can't be 0"): self.v // 0 def test_wrong_mod(self): with pytest.raises(ValueError, match="right part of expression can't be 0"): self.v % 0 @pytest.mark.parametrize("op, sign", ((operator.and_, "/\\"), (operator.or_, "\\/"), (operator.xor, "xor"))) def test_binary_logical(self, op, sign): result = op(self.p > self.v, 2 > self.v) assert to_str( result ) == f"(({self.p.name} > {self.v.name}) {sign} ({self.v.name} < 2))" def test_invert(self): result = ~(self.p < self.v) assert to_str(result) == f"(not ({self.p.name} < {self.v.name}))"
def test_var_array_without_shape(): with pytest.raises(ValueError, match="shape wasn't specified"): zn.Array(zn.var(int))
def _get_variable(seq) -> var: if is_range(seq): return var(int) assert False, f"{type(seq)} isn't supported, please use Sequence or Array here"
def test_float_start(): v = var(range(1.3, 10)) assert v.type.start == pytest.approx(1.3) assert v.type.stop == 10
def test_step_arg(): with pytest.raises(ValueError, match="Step values other than 1 are not supported"): v = var(range(10, 100, 2))
def test_one_arg(): v = var(range(100)) assert v.type.start == 0 assert v.type.stop == 100
class TestPos: def test_int_1d(self): a = zn.Array([1, 2, 3]) a = a[1] assert a.pos == (1, ) def test_tuple_1d(self): a = zn.Array([1, 2, 3]) with pytest.raises( ValueError, match="Array has 1 dimensions but 2 were specified"): _ = a[1, 0] def test_neg_int(self): a = zn.Array([1, 2, 3]) with pytest.raises(ValueError, match="Negative indexes are not supported for now"): _ = a[-1] def test_step(self): a = zn.Array([1, 2, 3]) with pytest.raises( ValueError, match="step other then 1 isn't supported, but it is 2"): _ = a[1:3:2] @pytest.mark.parametrize( "pos", [slice(-1, 2), slice(1, -2), slice(-1, -2)]) def test_neg_slice_2d(self, pos): a = zn.Array([[1], [2], [3]]) with pytest.raises(ValueError, match="Negative indexes are not supported for now"): _ = a[pos] @pytest.mark.parametrize("start", (10, 15)) def test_slice_wrong_start(self, start): array = zn.Array(zn.var(int), shape=(3, 2, 4)) with pytest.raises( ValueError, match=re.escape( f"start({start}) should be smaller then stop(10)")): _ = array[1, start:10] @pytest.mark.parametrize("pos, expected", [((0, 1), (0, 1)), ((0, slice(1, 3)), (0, slice(1, 3, 1)))]) def test_2d(self, pos, expected): a = zn.Array([[1, 2], [2, 3], [3, 4]]) a = a[pos] assert a.pos == expected @pytest.mark.parametrize("pos", (1, slice(1, 2))) @pytest.mark.parametrize("array", (zn.Array( [[1, 2], [2, 3], [3, 4]]), zn.Array(zn.var(int), shape=(3, 2, 4)))) def test_added_last_dim(self, pos, array): ndim = array.ndims() array = array[pos] assert len(array.pos) == ndim assert array.pos[1].start == 0 assert array.pos[1].stop.op == _Op_code.size assert array.pos[1].step == 1
def create_array(name, ndims): p = zn.Array(zn.var(range(5)), shape=(5,) * ndims) p._name = name return p
def test_float_stop(): v = var(range(1, 10.5)) assert v.type.start == 1 assert v.type.stop == pytest.approx(10.5)
def __init__(self): self.a = zn.var(range(100))
def __init__(self, a: int, b: int): super().__init__() self.a = par(a) self.b = par(b) self.x = var(range(-100, 101)) self.constraints = [self.a < self.x, self.x < self.b]
def test_float_start_and_stop(): v = var(range(3.14, 10.5)) assert v.type.start == pytest.approx(3.14) assert v.type.stop == pytest.approx(10.5)
def test_simple(): v = var(range(10, 100)) assert v.type.start == 10 assert v.type.stop == 100
def test_with_op(par_val): x = par(4) v = var(range(x, 10)) assert v.type.start == x assert v.type.stop == 10
def create_var(name): v = zn.var(int) v._name = name return v
def __init__(self): self.a = zn.var(range(3)) self.b = zn.var(range(3)) self.c = zn.var(range(3)) self.constraints = [constraint((self.a, self.b, self.c))]
def __init__(self, a): self.a = zn.par(a) self.x = zn.var(float) self.constraints = [self.x == self.a + 1]
def __init__(self, ): self.a = zn.Set(zn.var(zn.range(5))) self.b = zn.Set(zn.range(10))
def __init__(self, low, high): self.a = zn.Array(zn.var(range(1, 10)), shape=40) self.n = zn.var(range(low, high)) self.constraints = [zn.ndistinct(self.a) == self.n]
def __init__(self, array): self.a = zn.Array(array) with pytest.raises( ValueError, match="All elements of the array should be the same type"): MyModel(array) def test_var_array_without_shape(): with pytest.raises(ValueError, match="shape wasn't specified"): zn.Array(zn.var(int)) @pytest.mark.parametrize( "array", (zn.Array([[1], [2]]), zn.Array(zn.var(int), shape=(1, 2)))) @pytest.mark.parametrize("indexes", ((2, 0), (0, 2))) def test_array_index_error(array, indexes): with pytest.raises(IndexError): array[indexes] @pytest.mark.parametrize("array", (((1, 0, 1), (1, 3)), [[1, 3], [1, 1, 0]], (((1, 1), 1), ))) def test_different_length(array): class MyModel(zn.Model): def __init__(self, array): self.a = zn.Array(array) with pytest.raises( ValueError,