Use generator with ruamel.yaml

Question

I would like to have a bunch of generators in my config dict. So I tried this:

@yaml.register_class
class UniformDistribution:
    yaml_tag = '!uniform'

    @classmethod
    def from_yaml(cls, a, node):
        for x in node.value:
            if x[0].value == 'min':
                min_ = float(x[1].value)
            if x[0].value == 'max':
                max_ = float(x[1].value)
        def f():
            while True: 
                yield np.random.uniform(min_, max_)
        g = f()
        return g

However, the parser never returns because generators are used internally to resolve reference like &A and *A. Therefore, something like returning (g,) is a fairly simple workaround, but I would prefer a solution where I don't need the additional and very confusing index 0 term in next(config['position_generator'][0]).

Any Ideas?

score 0 · Accepted Answer · answered Mar 27 '20 at 17:36

This wrapper adapted from a different question did exactly what I was looking for.

class GeneratorWrapper(Generator):
    def __init__(self, function, *args):
        self.function = function
        self.args = args

    def send(self, ignored_arg):
        return self.function(*self.args)

    def throw(self, typ=None, val=None, tb=None):
        raise StopIteration


@yaml.register_class
class UniformDistribution:
    yaml_tag = '!uniform'

    @classmethod
    def from_yaml(cls, constructor, node):
        for x in node.value:
            value = float(x[1].value)
            if x[0].value == 'min':
                min_ = value
            if x[0].value == 'max':
                max_ = value
        return GeneratorWrapper(np.random.uniform, min_, max_)

Use generator with ruamel.yaml

1 Answers1