Source code for pyquchk.arbitraries.sequences

from math import log, ceil
from itertools import product
from . import Arbitrary, filter_possible
from .numbers import int_


@Arbitrary.set_for(list)
class list_(Arbitrary):
    """ List with length and elements taken from specified Arbitraries. """

    def __init__(self, length=int_(0, 32), elements=int_):
        self.length = length
        self.elements = elements

    def could_generate(self, x):
        return self.length.could_generate(len(x)) and \
               all(self.elements.could_generate(el) for el in x)

    def next_random(self):
        return [self.elements.next_random() for _ in range(self.length.next_random())]

    @filter_possible
    def gen_serial(self, amount):
        for l in self.length.gen_serial(amount):
            if l == 0:
                yield []
                continue

            el_amount = int(ceil((amount / 10) ** (1 / l)))
            elems = self.elements.gen_serial(el_amount)
            for res in product(elems, repeat=l):
                yield list(res)

    @filter_possible
    def shrink(self, x):
        l = len(x)
        if l > 1:
            yield x[:-1]
            yield x[1:]

            yield x[:l // 3]
            yield x[l // 3:2 * l // 3]
            yield x[2 * l // 3:]

            el_gens = map(self.elements.shrink, x)
            for els in zip(*el_gens):
                yield list(els)
        elif l == 1:
            for el in self.elements.shrink(x[0]):
                yield [el]