from .base import TohuBaseGenerator

__all__ = ['Apply', 'Lookup']


class DerivedGenerator(TohuBaseGenerator):
    """
    Base class for all derived generators
    """


class Apply(DerivedGenerator):

    def __init__(self, callable, *arg_gens, **kwarg_gens):
        super().__init__()
        self.callable = callable
        self.arg_gens_orig = arg_gens
        self.kwarg_gens_orig = kwarg_gens

        self.arg_gens = [g.clone() for g in self.arg_gens_orig]
        self.kwarg_gens = {name: g.clone() for name, g in self.kwarg_gens_orig.items()}
        self._input_generators = [g for g in self.arg_gens_orig] + [g for g in self.kwarg_gens_orig.values()]
        self._constituent_generators = [g for g in self.arg_gens] + [g for g in self.kwarg_gens.values()]

    def __next__(self):
        args = [next(g) for g in self.arg_gens]
        kwargs = {name: next(g) for name, g in self.kwarg_gens.items()}
        return self.callable(*args, **kwargs)

    def reset(self, seed):
        super().reset(seed)

    def spawn(self):
        return Apply(self.callable, *self.arg_gens, **self.kwarg_gens)


class Lookup(Apply):

    def __init__(self, g, mapping):
        self.g = g
        self.mapping = mapping

        def f_lookup(key):
            return mapping[key]

        super().__init__(f_lookup, g)

    def spawn(self):
        return Lookup(self.g, self.mapping)

@@ -1,8 +1,35 @@
+from random import Random%0A%0A
 from .ba
@@ -55,16 +55,68 @@
 enerator
+%0Afrom .primitive_generators import as_tohu_generator
 %0A%0A__all_
@@ -137,16 +137,34 @@
 'Lookup'
+, 'SelectMultiple'
 %5D%0A%0A%0Aclas
@@ -1477,8 +1477,640 @@
 apping)%0A
+%0A%0Aclass SelectMultiple(Apply):%0A%0A    def __init__(self, parent, num):%0A        parent = as_tohu_generator(parent)%0A        num = as_tohu_generator(num)%0A%0A        self.parent = parent%0A        self.num = num%0A        self.randgen = Random()%0A        func = self.randgen.choices%0A        super().__init__(func, parent, k=self.num)%0A%0A    def reset(self, seed):%0A        super().reset(seed)%0A        self.randgen.seed(seed)%0A%0A    def spawn(self):%0A        return SelectMultiple(self.parent, self.num)%0A%0A    def _set_random_state_from(self, other):%0A        super()._set_random_state_from(other)%0A        self.randgen.setstate(other.randgen.getstate())%0A

Add SelectMultiple (as derived generator)