from .base import TohuBaseGenerator

__all__ = ['Apply', 'DerivedGenerator']


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._constituents = [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_orig, **self.kwarg_gens_orig)
@@ -68,16 +68,26 @@
 nerator'
+, 'Lookup'
 %5D%0A%0A%0Aclas
@@ -1014,8 +1014,283 @@
 ns_orig)
+%0A%0A%0Aclass Lookup(Apply):%0A%0A    def __init__(self, g, mapping):%0A        self.g = g%0A        self.mapping = mapping%0A%0A        def f_lookup(key):%0A            return mapping%5Bkey%5D%0A%0A        super().__init__(f_lookup, g)%0A%0A    def spawn(self):%0A        return Lookup(self.g, self.mapping)

Add Lookup