Source code for pyquchk.arbitraries.arbitrary

import six
from itertools import islice, count
from functools import wraps

registry = {}
registry[None] = []


class combomethod(object):
    def __init__(self, method):
        self.method = method

    def __get__(self, obj=None, objtype=None):
        @wraps(self.method)
        def _wrapper(*args, **kwargs):
            if obj is not None:
                return self.method(obj, *args, **kwargs)
            else:
                return self.method(objtype.inst, *args, **kwargs)

        return _wrapper


class ArbitraryMeta(type):
    arbitrary_cls = None

    def __init__(cls, name, bases, attrs):
        super(ArbitraryMeta, cls).__init__(name, bases, attrs)

    def __new__(meta, name, bases, attrs):
        cls = super(ArbitraryMeta, meta).__new__(meta, name, bases, attrs)

        if name == 'Arbitrary':
            # set the Arbitrary class to an ArbitraryMeta class variable
            # we can't just access Arbitrary class in this __new__ method
            meta.arbitrary_cls = cls
            return cls
        if not meta.arbitrary_cls:
            # Arbitrary class isn't created yet, now creating its superclass
            return cls

        # Arbitrary class already created, now creating some its subclass

        # add hard limiting to gen_serial method
        gen_serial_old = getattr(cls, 'gen_serial')

        @wraps(gen_serial_old)
        def gen_serial_new(self, amount, gen_serial=gen_serial_old):
            return islice(gen_serial(self, amount), amount)

        setattr(cls, 'gen_serial', gen_serial_new)

        for meth_name in ['next_random', 'gen_serial', 'could_generate',
                          'shrink']:
            meth_src = getattr(cls, meth_name)
            meth_new = combomethod(meth_src)
            setattr(cls, meth_name, meth_new)

        cls.inst = cls()
        registry[None].append(cls)
        return cls

    def __call__(cls, *args, **kwargs):
        return super(ArbitraryMeta, cls).__call__(*args, **kwargs)


class Arbitrary(six.with_metaclass(ArbitraryMeta)):
    def next_random(self):
        """Get next random value.

        You should override this method."""
        raise StopIteration

    def gen_serial(self, amount):
        """Generate or return a sequence of serial values (finite or infinite amount).

        You should override this method."""
        return []

    def could_generate(self, x):
        """Check if x could've been generated by self.

        You can override this method, but it's not used now in any way."""
        raise NotImplementedError

    def shrink(self, x):
        """Get possible shrinkings of x.
        You can assume here that self.could_generate(x) == True.

        You have to override this method if you want shrinking to be available."""
        return []

    @classmethod
    def sample(cls):
        try:
            args = cls.__init__.args_for_sample['args']
            kwargs = cls.__init__.args_for_sample['kwargs']
        except AttributeError:
            args = []
            kwargs = {}

        inst = cls(*args, **kwargs)

        dct = {
            'args': {
                'args': args,
                'kwargs': kwargs,
            },
            'random': (inst.next_random() for _ in count()),
            'serial': (
                (2 ** p2, inst.gen_serial(2 ** p2))
                for p2 in count()
            )
        }
        return dct

    @staticmethod
    def args_for_sample(*args, **kwargs):
        def wrapper(method):
            method.args_for_sample = {'args': args, 'kwargs': kwargs}
            return method

        return wrapper

    @staticmethod
    def set_for(type):
        def wrapper(arb):
            if type in registry:
                raise LookupError('Arbitrary class already set for %s' % type)
            registry[type] = arb
            return arb

        return wrapper

    @staticmethod
    def get_for(type):
        """Get default Arbitrary for the specified type."""
        try:
            return registry[type]
        except KeyError:
            raise LookupError('No Arbitrary class is set for %s' % type)

    @staticmethod
    def get_all():
        return registry[None]