CODEWITHSUNDEEP
python
Yash
Yash

Reputation: 1008

Overriding a method that yields

I am new to Python 3 concepts.

The base class has following functions:

class DocumentFormatter(object):
    def transform_element(self, key, value):
        if isinstance(value, list):
            for li, lv in enumerate(value):
                for inner_k, inner_v in self.transform_element(
                        "%s.%s" % (key, li), lv):
                    yield inner_k, inner_v
        elif isinstance(value, dict):
            formatted = self.format_document(value)
            for doc_key in formatted:
                yield "%s.%s" % (key, doc_key), formatted[doc_key]
        else:
            # We assume that transform_value will return a 'flat' value,
            # not a list or dict
            yield key, self.transform_value(value)

    def format_document(self, document):
        def flatten(doc, path):
            top_level = (len(path) == 0)
            if not top_level:
                path_string = ".".join(path)
            for k in doc:
                v = doc[k]
                if isinstance(v, dict):
                    path.append(k)
                    for inner_k, inner_v in flatten(v, path):
                        yield inner_k, inner_v
                    path.pop()
                else:
                    transformed = self.transform_element(k, v)
                    for new_k, new_v in transformed:
                        if top_level:
                            yield new_k, new_v
                        else:
                            yield "%s.%s" % (path_string, new_k), new_v
        return dict(flatten(document, []))

In the derived class, I only want to change a bit of the transform_element function:

def transform_element(self, key, value):
    if key=="cats":
        yield key, self.transform_value(value)
    else:
        yield super().transform_element(key,value)

If the key is "cats", I want to yield using my logic. Else, I want the base class implementation to work. Am I supposed to call yield or should I call return?

Upvotes: 1

Views: 544

Answers (1)

dopstar
dopstar

Reputation: 1488

you should iterate over the overridden method and yield each item of it, otherwise you will end up with a generator of a generator which is not what you need.

So your method should look like:

def transform_element(self, key, value):
    if key=="cats":
        yield key, self.transform_value(value)
    else:
        for item in super().transform_element(key,value):
            yield item

Upvotes: 2

Related Questions