def test_map_with_lookahead(): def if_one(x): return 'one' + str(x) def if_many(x): return 'many' + str(x) assert list(map_with_lookahead(iter([]), if_one, if_many)) == [] assert list(map_with_lookahead(iter([1]), if_one, if_many)) == [if_one(1)] assert list(map_with_lookahead(range(5), if_one, if_many)) == list(map(if_many, range(5))) assert list(map_with_lookahead(range(10), if_one=if_one)) == list(range(10)) assert list(map_with_lookahead(iter([1]), if_many=if_many)) == [1]
def process_file(path): docs = load_documents(path) if not uri: for doc in docs: yield path, doc else: url = as_url(path) def add_uri_no_part(x): idx, doc = x return url, doc def add_uri_with_part(x): idx, doc = x return mk_part_uri(url, idx), doc yield from map_with_lookahead(enumerate(docs), if_one=add_uri_no_part, if_many=add_uri_with_part)