-
-
Notifications
You must be signed in to change notification settings - Fork 106
Expand file tree
/
Copy pathtest_parallel_imap.py
More file actions
121 lines (95 loc) · 3.6 KB
/
Copy pathtest_parallel_imap.py
File metadata and controls
121 lines (95 loc) · 3.6 KB
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
"""Lazy ordered `imap` and completion-order `imap_unordered`."""
from __future__ import annotations
import contextlib
import operator
import threading
import time
import typing
import pytest
from progressbar._parallel import _sync
_executed: set[int] = set()
_executed_lock: threading.Lock = threading.Lock()
def _double(value: int) -> int:
return value * 2
def _record(value: int) -> int:
with _executed_lock:
_executed.add(value)
return value
def _sleep_inverse(value: int) -> int:
# Later items finish first, scrambling completion order.
time.sleep((5 - value) * 0.03)
return value
def _boom_on_two(value: int) -> int:
if value == 2:
raise ValueError('boom')
return value
class TestImap:
def test_yields_results_in_input_order(self) -> None:
assert list(_sync.imap(_double, range(10), bar=False)) == [
value * 2 for value in range(10)
]
@pytest.mark.no_freezegun
def test_ordered_despite_scrambled_completion(self) -> None:
assert list(
_sync.imap(_sleep_inverse, range(5), workers=5, bar=False)
) == list(range(5))
def test_lazy(self) -> None:
iterator: typing.Generator[typing.Any, None, None] = _sync.imap(
_double, range(10), workers=1, bar=False
)
assert next(iterator) == 0
iterator.close()
def test_on_error_return_yields_exceptions_in_place(self) -> None:
results: list[typing.Any] = list(
_sync.imap(_boom_on_two, range(4), on_error='return', bar=False)
)
assert results[0] == 0
assert isinstance(results[2], ValueError)
assert results[3] == 3
def test_on_error_raise_raises_at_iteration(self) -> None:
iterator = _sync.imap(_boom_on_two, range(4), workers=1, bar=False)
with pytest.raises(ValueError, match='boom'):
list(iterator)
class TestImapUnordered:
@pytest.mark.no_freezegun
def test_yields_pairs_in_completion_order(self) -> None:
pairs: list[tuple[int, int]] = list(
_sync.imap_unordered(
_sleep_inverse, range(5), workers=5, bar=False
)
)
assert sorted(pairs) == [(value, value) for value in range(5)]
# The fastest item (highest input) completes and is seen first.
assert pairs[0] == (4, 4)
def test_multi_iterable_pairs_use_args_tuple(self) -> None:
pairs = list(
_sync.imap_unordered(
operator.add, [1, 2], [10, 20], workers=1, bar=False
)
)
assert sorted(pairs) == [((1, 10), 11), ((2, 20), 22)]
def test_on_error_return_pairs_exceptions(self) -> None:
pairs = dict(
_sync.imap_unordered(
_boom_on_two, range(4), on_error='return', bar=False
)
)
assert isinstance(pairs[2], ValueError)
assert pairs[3] == 3
class TestGeneratorCleanup:
def test_early_break_stops_submission(self) -> None:
_executed.clear()
for item, _result in _sync.imap_unordered(
_record, range(100), workers=1, buffersize=2, bar=False
):
if item >= 1:
break
# workers=1, window=2: only the in-window items ever ran; the
# other ~97 must have been cancelled by the generator close.
time.sleep(0.2) # let any straggler drain before asserting
assert len(_executed) <= 5
def test_contextlib_closing_recipe(self) -> None:
with contextlib.closing(
_sync.imap(_double, range(10), workers=1, bar=False)
) as iterator:
assert next(iterator) == 0