Source code for pyquchk.arbitraries.arbitrary

"""Contains :class:`Arbitrary` class and the machinery it needs to work.
You shouldn't need anything other than this class from here.
"""
from abc import ABCMeta, abstractmethod
import re
import inspect
import six
from itertools import islice, count
from warnings import warn
from .utils import until_possible, filter_possible

registry = {}


class Arbitrary(six.with_metaclass(ABCMeta)):
    """Base class for all arbitraries, both built-in and custom.

    To create a custom arbitrary, you need to subclass from :class:`Arbitrary`
    and override (some or all) instance methods listed below."""

    @abstractmethod
    def next_random(self):
        """Get next random value."""
        raise StopIteration

    @abstractmethod
    def gen_serial(self):
        """Generate or return a sequence (finite or infinite) of serial values.
        """
        return []

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

        Required iff you use decorators like :func:`.filter_possible` or
        :func:`.until_possible`."""
        raise NotImplementedError

    @abstractmethod
    def shrink(self, x):
        """Get possible shrinkings of ``x``.

        :param x: Value previously generated by this arbitrary.
        """
        return []

    def gen_random(self):
        while True:
            yield self.next_random()

    @classmethod
    def samples_docstring_part(cls):
        try:
            args_formatted = cls.__init__.args_for_sample['def']
        except AttributeError:
            args_formatted = ''

        def append(line, indent=' ' * 4):
            res.append(indent + line)

        res = []
        append('', indent='')
        append('', indent='')
        append('.. rst-class:: html-toggle', indent='')
        append('', indent='')
        append('**Sample values**', indent='')
        append('', indent='')
        append('.. ipython::', indent='')
        append('')
        append('>>> arb = %s(%s)' % (cls.__name__, args_formatted))
        append('')
        append('@notest')
        append('>>> list(islice(arb.gen_random(), 20))')
        append('')
        append('@notest')
        append('>>> list(islice(arb.gen_serial(), 50))')
        append('')
        return '\n'.join(res)

    @staticmethod
    def args_for_sample(*args, **kwargs):
        def wrapper(method):
            method_src = inspect.getsource(method)
            m = re.match(r'\s*@Arbitrary.args_for_sample\((.*)\)', method_src)
            argsdef = m.group(1)
            method.args_for_sample = {
                'args': args,
                'kwargs': kwargs,
                'def': argsdef
            }
            return method

        return wrapper

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

        if arb is None:
            return wrapper
        else:
            return wrapper(arb)

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


def with_samples(arb_cls):
    arb_cls.__doc__ += arb_cls.samples_docstring_part()
    return arb_cls


spec_processors = []


def get_arbitrary(spec):
    if isinstance(spec, Arbitrary):
        return spec
    elif isinstance(spec, type):
        if issubclass(spec, Arbitrary):
            return spec()
        try:
            return registry[spec]()
        except KeyError:
            raise LookupError('No Arbitrary class is set for %s' % spec)

    spec_processed = [proc(spec) for proc in spec_processors]
    spec_processed = [pr for pr in spec_processed if pr is not None]
    if not spec_processed:
        raise LookupError('No match found for specification: %r' % spec)
    elif len(spec_processed) > 1:
        raise LookupError('Ambigous specification: %r' % spec)
    else:
        return spec_processed[0]


def add_spec_processor(processor):
    spec_processors.append(processor)
    return processor


class MappedArbitrary(Arbitrary):
    @abstractmethod
    def pack(self, x):
        pass

    @abstractmethod
    def unpack(self, x):
        pass

    def next_random(self):
        return self.pack(self.inner_arb.next_random())

    def gen_serial(self):
        for el in self.inner_arb.gen_serial():
            yield self.pack(el)

    def could_generate(self, x):
        return self.inner_arb.could_generate(self.unpack(x))

    def shrink(self, x):
        return map(self.pack, self.inner_arb.shrink(self.unpack(x)))


class FilteredArbitrary(Arbitrary):
    @abstractmethod
    def predicate(self, x):
        pass

    def could_generate(self, x):
        return self.inner_arb.could_generate(x) and \
               self.predicate(x)

    @until_possible
    def next_random(self):
        return self.inner_arb.next_random()

    @filter_possible
    def gen_serial(self):
        return self.inner_arb.gen_serial()

    @filter_possible
    def shrink(self, x):
        return self.inner_arb.shrink(x)