For subscriptable iterable that implements __getitem__, save range then return a generator that iterates given range.
def split_on_condition_gen(seq, condition):
def inner_gen(range_: range):
for n in range_:
yield seq[n]
last = 0
current = 0
for item in seq:
if condition(item):
current += 1
else:
yield inner_gen(range(last, current))
last = current
As it checks until condition breaks, it's not completely lazy, but at least it can handle large input as it won't consume extra memory until you really need those.
>>> from this import s as seq_
~~ ZEN OF PYTHON GOES HERE ~~
>>> condition = lambda x: x != ','
>>> output = split_on_condition_gen(seq_, condition)
>>> output
<generator object split_on_condition_gen at 0x000002602C1ECEB0>
>>> list_output = list(output)
>>> list_output[2]
<generator object split_on_condition_gen.<locals>.inner_gen at 0x000002602C1ECEB0>
>>> list(list_output[2])
['l', ',', ' ', 'e', 'r', 's', 'h', 'f', 'r', ' ', ...]
However, for iterables without __getitem__ implemented - for ones like this:
class NonIndexAble:
def __init__(self, initial: str):
self.source = initial
self.list_ = list(initial)
def __iter__(self):
return self
def __next__(self):
try:
return self.list_.pop(0)
except IndexError:
raise StopIteration()
Will raise errors like this:
>>> non_index_able = NonIndexAble('aaabaaaaba')
>>> condition = lambda x: x == "b"
>>> for part in split_on_condition_gen(non_index_able, condition):
... for item in part:
... pass
Traceback (most recent call last):
File "<input>", line 2, in <module>
File "<input>", line 11, in inner_gen
TypeError: 'NonIndexAble' object is not subscriptable
First I thought we could use itertools.islice, but found out that it consumes iterator passed, shown like this:
except StopIteration:
# Consume to *stop*.
for i, element in zip(range(i + 1, stop), iterable):
pass
Rather, as most use-case of this generator would be nested for loop, it will run inner-generators in sequence they've created.
For that case, we could just have single generator instance and have counter variable for 'cursor' location. Then we can get values by skipping from counter to start and yield from start to end index and leave generator as-is, not consuming completely like islice does so next 'block' could continue with same generator instance.
counter = 0
def inner_gen(seq_, start, end):
nonlocal counter
for _ in range(counter, start): # Skip from counter to start
next(seq_)
for _ in range(start, end): # yield from start to end
yield next(seq_)
counter = end
Full implementation:
def split_on_condition_gen(seq, condition):
seq_copy, seq_current = tee(seq)
counter = 0
def inner_gen(seq_, start, end):
nonlocal counter
for _ in range(counter, start): # Skip from counter to start
next(seq_)
for _ in range(start, end): # yield from start to end
yield next(seq_)
counter = end
last = 0
for idx, item in enumerate(seq_current):
if condition(item):
continue
yield inner_gen(seq_copy, last, idx)
last = idx + 1
yield inner_gen(seq_copy, last, idx + 1)
Test:
>>> non_index_able = NonIndexAble('asdf adsfa ds sdfdf adf')
>>> condition = lambda x: x != ' '
>>> for part in split_on_condition_gen(non_index_able, condition):
... print("Part start")
... for item in part:
... print(item, end=' ')
... print("\nEnd")
Part start
a s d f
End
Part start
a d s f a
End
Part start
d s
End
Part start
s d f d f
End
Part start
a d f
End
