GridTools/gt4py

Cleanup duplication of flatten of sequences

havogt opened this issue · 1 comments

havogt commented

We have many variations of flatten implementations in our code base, we should provide a single (customizable) one for all use-cases.

This is the result of research in #1489 (with reference impl there):

There already are implementations for flatten in eve.utils. They are very general and treat any iterable as something to flatten. One frequent usage that seems to come up is flattening iterables that may contain TupleType (which may contain itself).

If we simply make TupleType iterable like so:

class TupleType(...):
    ...
    def __iter__(self) -> Iterator[DataType]:
        yield from self.types

We can use eve.utils.flatten_iter to achieve most of the desired effects:

nested_types = [
    (type1, type2),
    type3,
    TupleType(types=[type4, type5, TupleType(...)]),
    ...
]

flattened_tuple = tuple(eve.utils.flatten_iter(nested_types))
flattened_list = list(eve_utils.flatten_iter(nested_types))
filtering_comprehension = [t for t in eve_utils.flatten_iter(nested_types) if t ... ]

The same may be possible in other use cases.