1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
|
"""Selection operators."""
from __future__ import annotations
import asyncio
import builtins
import collections
from typing import Awaitable, Callable, TypeVar, AsyncIterable, AsyncIterator
from . import transform
from ..aiter_utils import aiter, anext
from ..core import streamcontext, pipable_operator
__all__ = [
"take",
"takelast",
"skip",
"skiplast",
"getitem",
"filter",
"until",
"dropwhile",
"takewhile",
]
T = TypeVar("T")
@pipable_operator
async def take(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
"""Forward the first ``n`` elements from an asynchronous sequence.
If ``n`` is negative, it simply terminates before iterating the source.
"""
enumerated = transform.enumerate.raw(source)
async with streamcontext(enumerated) as streamer:
if n <= 0:
return
async for i, item in streamer:
yield item
if i >= n - 1:
return
@pipable_operator
async def takelast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
"""Forward the last ``n`` elements from an asynchronous sequence.
If ``n`` is negative, it simply terminates after iterating the source.
Note: it is required to reach the end of the source before the first
element is generated.
"""
queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
async with streamcontext(source) as streamer:
async for item in streamer:
queue.append(item)
for item in queue:
yield item
@pipable_operator
async def skip(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
"""Forward an asynchronous sequence, skipping the first ``n`` elements.
If ``n`` is negative, no elements are skipped.
"""
enumerated = transform.enumerate.raw(source)
async with streamcontext(enumerated) as streamer:
async for i, item in streamer:
if i >= n:
yield item
@pipable_operator
async def skiplast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
"""Forward an asynchronous sequence, skipping the last ``n`` elements.
If ``n`` is negative, no elements are skipped.
Note: it is required to reach the ``n+1`` th element of the source
before the first element is generated.
"""
queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
async with streamcontext(source) as streamer:
async for item in streamer:
if n <= 0:
yield item
continue
if len(queue) == n:
yield queue[0]
queue.append(item)
@pipable_operator
async def filterindex(
source: AsyncIterable[T], func: Callable[[int], bool]
) -> AsyncIterator[T]:
"""Filter an asynchronous sequence using the index of the elements.
The given function is synchronous, takes the index as an argument,
and returns ``True`` if the corresponding should be forwarded,
``False`` otherwise.
"""
enumerated = transform.enumerate.raw(source)
async with streamcontext(enumerated) as streamer:
async for i, item in streamer:
if func(i):
yield item
@pipable_operator
def slice(source: AsyncIterable[T], *args: int) -> AsyncIterator[T]:
"""Slice an asynchronous sequence.
The arguments are the same as the builtin type slice.
There are two limitations compare to regular slices:
- Positive stop index with negative start index is not supported
- Negative step is not supported
"""
s = builtins.slice(*args)
start, stop, step = s.start or 0, s.stop, s.step or 1
aiterator = aiter(source)
# Filter the first items
if start < 0:
aiterator = takelast.raw(aiterator, abs(start))
elif start > 0:
aiterator = skip.raw(aiterator, start)
# Filter the last items
if stop is not None:
if stop >= 0 and start < 0:
raise ValueError("Positive stop with negative start is not supported")
elif stop >= 0:
aiterator = take.raw(aiterator, stop - start)
else:
aiterator = skiplast.raw(aiterator, abs(stop))
# Filter step items
if step is not None:
if step > 1:
aiterator = filterindex.raw(aiterator, lambda i: i % step == 0)
elif step < 0:
raise ValueError("Negative step not supported")
# Return
return aiterator
@pipable_operator
async def item(source: AsyncIterable[T], index: int) -> AsyncIterator[T]:
"""Forward the ``n``th element of an asynchronous sequence.
The index can be negative and works like regular indexing.
If the index is out of range, and ``IndexError`` is raised.
"""
# Prepare
if index >= 0:
source = skip.raw(source, index)
else:
source = takelast(source, abs(index))
async with streamcontext(source) as streamer:
# Get first item
try:
result = await anext(streamer)
except StopAsyncIteration:
raise IndexError("Index out of range")
# Check length
if index < 0:
count = 1
async for _ in streamer:
count += 1
if count != abs(index):
raise IndexError("Index out of range")
# Yield result
yield result
@pipable_operator
def getitem(source: AsyncIterable[T], index: int | builtins.slice) -> AsyncIterator[T]:
"""Forward one or several items from an asynchronous sequence.
The argument can either be a slice or an integer.
See the slice and item operators for more information.
"""
if isinstance(index, builtins.slice):
return slice.raw(source, index.start, index.stop, index.step)
if isinstance(index, int):
return item.raw(source, index)
raise TypeError("Not a valid index (int or slice)")
@pipable_operator
async def filter(
source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
"""Filter an asynchronous sequence using an arbitrary function.
The function takes the item as an argument and returns ``True``
if it should be forwarded, ``False`` otherwise.
The function can either be synchronous or asynchronous.
"""
iscorofunc = asyncio.iscoroutinefunction(func)
async with streamcontext(source) as streamer:
async for item in streamer:
result = func(item)
if iscorofunc:
assert isinstance(result, Awaitable)
result = await result
if result:
yield item
@pipable_operator
async def until(
source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
"""Forward an asynchronous sequence until a condition is met.
Contrary to the ``takewhile`` operator, the last tested element is included
in the sequence.
The given function takes the item as an argument and returns a boolean
corresponding to the condition to meet. The function can either be
synchronous or asynchronous.
"""
iscorofunc = asyncio.iscoroutinefunction(func)
async with streamcontext(source) as streamer:
async for item in streamer:
result = func(item)
if iscorofunc:
assert isinstance(result, Awaitable)
result = await result
yield item
if result:
return
@pipable_operator
async def takewhile(
source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
"""Forward an asynchronous sequence while a condition is met.
Contrary to the ``until`` operator, the last tested element is not included
in the sequence.
The given function takes the item as an argument and returns a boolean
corresponding to the condition to meet. The function can either be
synchronous or asynchronous.
"""
iscorofunc = asyncio.iscoroutinefunction(func)
async with streamcontext(source) as streamer:
async for item in streamer:
result = func(item)
if iscorofunc:
assert isinstance(result, Awaitable)
result = await result
if not result:
return
yield item
@pipable_operator
async def dropwhile(
source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
"""Discard the elements from an asynchronous sequence
while a condition is met.
The given function takes the item as an argument and returns a boolean
corresponding to the condition to meet. The function can either be
synchronous or asynchronous.
"""
iscorofunc = asyncio.iscoroutinefunction(func)
async with streamcontext(source) as streamer:
async for item in streamer:
result = func(item)
if iscorofunc:
assert isinstance(result, Awaitable)
result = await result
if not result:
yield item
break
async for item in streamer:
yield item
|
