Skip to content

Commit f537b2a

Browse files
asvetlovgvanrossum
asvetlov
and
gvanrossum
authored
bpo-46771: Implement asyncio context managers for handling timeouts (GH-31394)
Example: async with asyncio.timeout(5): await some_task() Will interrupt the await and raise TimeoutError if some_task() takes longer than 5 seconds. Co-authored-by: Guido van Rossum <guido@python.org>
1 parent 32bf359 commit f537b2a

File tree

4 files changed

+384
-0
lines changed

4 files changed

+384
-0
lines changed

‎Lib/asyncio/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from .subprocessimport*
1919
from .tasksimport*
2020
from .taskgroupsimport*
21+
from .timeoutsimport*
2122
from .threadsimport*
2223
from .transportsimport*
2324

@@ -34,6 +35,7 @@
3435
subprocess.__all__+
3536
tasks.__all__+
3637
threads.__all__+
38+
timeouts.__all__+
3739
transports.__all__)
3840

3941
ifsys.platform=='win32': # pragma: no cover

‎Lib/asyncio/timeouts.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
importenum
2+
3+
fromtypesimportTracebackType
4+
fromtypingimportfinal, Optional, Type
5+
6+
from . importevents
7+
from . importexceptions
8+
from . importtasks
9+
10+
11+
__all__= (
12+
"Timeout",
13+
"timeout",
14+
"timeout_at",
15+
)
16+
17+
18+
class_State(enum.Enum):
19+
CREATED="created"
20+
ENTERED="active"
21+
EXPIRING="expiring"
22+
EXPIRED="expired"
23+
EXITED="finished"
24+
25+
26+
@final
27+
classTimeout:
28+
29+
def__init__(self, when: Optional[float]) ->None:
30+
self._state=_State.CREATED
31+
32+
self._timeout_handler: Optional[events.TimerHandle] =None
33+
self._task: Optional[tasks.Task] =None
34+
self._when=when
35+
36+
defwhen(self) ->Optional[float]:
37+
returnself._when
38+
39+
defreschedule(self, when: Optional[float]) ->None:
40+
assertself._stateisnot_State.CREATED
41+
ifself._stateisnot_State.ENTERED:
42+
raiseRuntimeError(
43+
f"Cannot change state of {self._state.value} Timeout",
44+
)
45+
46+
self._when=when
47+
48+
ifself._timeout_handlerisnotNone:
49+
self._timeout_handler.cancel()
50+
51+
ifwhenisNone:
52+
self._timeout_handler=None
53+
else:
54+
loop=events.get_running_loop()
55+
self._timeout_handler=loop.call_at(
56+
when,
57+
self._on_timeout,
58+
)
59+
60+
defexpired(self) ->bool:
61+
"""Is timeout expired during execution?"""
62+
returnself._statein (_State.EXPIRING, _State.EXPIRED)
63+
64+
def__repr__(self) ->str:
65+
info= ['']
66+
ifself._stateis_State.ENTERED:
67+
when=round(self._when, 3) ifself._whenisnotNoneelseNone
68+
info.append(f"when={when}")
69+
info_str=' '.join(info)
70+
returnf"<Timeout [{self._state.value}]{info_str}>"
71+
72+
asyncdef__aenter__(self) ->"Timeout":
73+
self._state=_State.ENTERED
74+
self._task=tasks.current_task()
75+
ifself._taskisNone:
76+
raiseRuntimeError("Timeout should be used inside a task")
77+
self.reschedule(self._when)
78+
returnself
79+
80+
asyncdef__aexit__(
81+
self,
82+
exc_type: Optional[Type[BaseException]],
83+
exc_val: Optional[BaseException],
84+
exc_tb: Optional[TracebackType],
85+
) ->Optional[bool]:
86+
assertself._statein (_State.ENTERED, _State.EXPIRING)
87+
88+
ifself._timeout_handlerisnotNone:
89+
self._timeout_handler.cancel()
90+
self._timeout_handler=None
91+
92+
ifself._stateis_State.EXPIRING:
93+
self._state=_State.EXPIRED
94+
95+
ifself._task.uncancel() ==0andexc_typeisexceptions.CancelledError:
96+
# Since there are no outstanding cancel requests, we're
97+
# handling this.
98+
raiseTimeoutError
99+
elifself._stateis_State.ENTERED:
100+
self._state=_State.EXITED
101+
102+
returnNone
103+
104+
def_on_timeout(self) ->None:
105+
assertself._stateis_State.ENTERED
106+
self._task.cancel()
107+
self._state=_State.EXPIRING
108+
# drop the reference early
109+
self._timeout_handler=None
110+
111+
112+
deftimeout(delay: Optional[float]) ->Timeout:
113+
"""Timeout async context manager.
114+
115+
Useful in cases when you want to apply timeout logic around block
116+
of code or in cases when asyncio.wait_for is not suitable. For example:
117+
118+
>>> async with asyncio.timeout(10): # 10 seconds timeout
119+
... await long_running_task()
120+
121+
122+
delay - value in seconds or None to disable timeout logic
123+
124+
long_running_task() is interrupted by raising asyncio.CancelledError,
125+
the top-most affected timeout() context manager converts CancelledError
126+
into TimeoutError.
127+
"""
128+
loop=events.get_running_loop()
129+
returnTimeout(loop.time() +delayifdelayisnotNoneelseNone)
130+
131+
132+
deftimeout_at(when: Optional[float]) ->Timeout:
133+
"""Schedule the timeout at absolute time.
134+
135+
Like timeout() but argument gives absolute time in the same clock system
136+
as loop.time().
137+
138+
Please note: it is not POSIX time but a time with
139+
undefined starting base, e.g. the time of the system power on.
140+
141+
>>> async with asyncio.timeout_at(loop.time() + 10):
142+
... await long_running_task()
143+
144+
145+
when - a deadline when timeout occurs or None to disable timeout logic
146+
147+
long_running_task() is interrupted by raising asyncio.CancelledError,
148+
the top-most affected timeout() context manager converts CancelledError
149+
into TimeoutError.
150+
"""
151+
returnTimeout(when)

‎Lib/test/test_asyncio/test_timeouts.py‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
"""Tests for asyncio/timeouts.py"""
2+
3+
importunittest
4+
importtime
5+
6+
importasyncio
7+
fromasyncioimporttasks
8+
9+
10+
deftearDownModule():
11+
asyncio.set_event_loop_policy(None)
12+
13+
14+
classTimeoutTests(unittest.IsolatedAsyncioTestCase):
15+
16+
asyncdeftest_timeout_basic(self):
17+
withself.assertRaises(TimeoutError):
18+
asyncwithasyncio.timeout(0.01) ascm:
19+
awaitasyncio.sleep(10)
20+
self.assertTrue(cm.expired())
21+
22+
asyncdeftest_timeout_at_basic(self):
23+
loop=asyncio.get_running_loop()
24+
25+
withself.assertRaises(TimeoutError):
26+
deadline=loop.time() +0.01
27+
asyncwithasyncio.timeout_at(deadline) ascm:
28+
awaitasyncio.sleep(10)
29+
self.assertTrue(cm.expired())
30+
self.assertEqual(deadline, cm.when())
31+
32+
asyncdeftest_nested_timeouts(self):
33+
loop=asyncio.get_running_loop()
34+
cancelled=False
35+
withself.assertRaises(TimeoutError):
36+
deadline=loop.time() +0.01
37+
asyncwithasyncio.timeout_at(deadline) ascm1:
38+
# Only the topmost context manager should raise TimeoutError
39+
try:
40+
asyncwithasyncio.timeout_at(deadline) ascm2:
41+
awaitasyncio.sleep(10)
42+
exceptasyncio.CancelledError:
43+
cancelled=True
44+
raise
45+
self.assertTrue(cancelled)
46+
self.assertTrue(cm1.expired())
47+
self.assertTrue(cm2.expired())
48+
49+
asyncdeftest_waiter_cancelled(self):
50+
loop=asyncio.get_running_loop()
51+
cancelled=False
52+
withself.assertRaises(TimeoutError):
53+
asyncwithasyncio.timeout(0.01):
54+
try:
55+
awaitasyncio.sleep(10)
56+
exceptasyncio.CancelledError:
57+
cancelled=True
58+
raise
59+
self.assertTrue(cancelled)
60+
61+
asyncdeftest_timeout_not_called(self):
62+
loop=asyncio.get_running_loop()
63+
t0=loop.time()
64+
asyncwithasyncio.timeout(10) ascm:
65+
awaitasyncio.sleep(0.01)
66+
t1=loop.time()
67+
68+
self.assertFalse(cm.expired())
69+
# 2 sec for slow CI boxes
70+
self.assertLess(t1-t0, 2)
71+
self.assertGreater(cm.when(), t1)
72+
73+
asyncdeftest_timeout_disabled(self):
74+
loop=asyncio.get_running_loop()
75+
t0=loop.time()
76+
asyncwithasyncio.timeout(None) ascm:
77+
awaitasyncio.sleep(0.01)
78+
t1=loop.time()
79+
80+
self.assertFalse(cm.expired())
81+
self.assertIsNone(cm.when())
82+
# 2 sec for slow CI boxes
83+
self.assertLess(t1-t0, 2)
84+
85+
asyncdeftest_timeout_at_disabled(self):
86+
loop=asyncio.get_running_loop()
87+
t0=loop.time()
88+
asyncwithasyncio.timeout_at(None) ascm:
89+
awaitasyncio.sleep(0.01)
90+
t1=loop.time()
91+
92+
self.assertFalse(cm.expired())
93+
self.assertIsNone(cm.when())
94+
# 2 sec for slow CI boxes
95+
self.assertLess(t1-t0, 2)
96+
97+
asyncdeftest_timeout_zero(self):
98+
loop=asyncio.get_running_loop()
99+
t0=loop.time()
100+
withself.assertRaises(TimeoutError):
101+
asyncwithasyncio.timeout(0) ascm:
102+
awaitasyncio.sleep(10)
103+
t1=loop.time()
104+
self.assertTrue(cm.expired())
105+
# 2 sec for slow CI boxes
106+
self.assertLess(t1-t0, 2)
107+
self.assertTrue(t0<=cm.when() <=t1)
108+
109+
asyncdeftest_foreign_exception_passed(self):
110+
withself.assertRaises(KeyError):
111+
asyncwithasyncio.timeout(0.01) ascm:
112+
raiseKeyError
113+
self.assertFalse(cm.expired())
114+
115+
asyncdeftest_foreign_exception_on_timeout(self):
116+
asyncdefcrash():
117+
try:
118+
awaitasyncio.sleep(1)
119+
finally:
120+
1/0
121+
withself.assertRaises(ZeroDivisionError):
122+
asyncwithasyncio.timeout(0.01):
123+
awaitcrash()
124+
125+
asyncdeftest_foreign_cancel_doesnt_timeout_if_not_expired(self):
126+
withself.assertRaises(asyncio.CancelledError):
127+
asyncwithasyncio.timeout(10) ascm:
128+
asyncio.current_task().cancel()
129+
awaitasyncio.sleep(10)
130+
self.assertFalse(cm.expired())
131+
132+
asyncdeftest_outer_task_is_not_cancelled(self):
133+
asyncdefouter() ->None:
134+
withself.assertRaises(TimeoutError):
135+
asyncwithasyncio.timeout(0.001):
136+
awaitasyncio.sleep(10)
137+
138+
task=asyncio.create_task(outer())
139+
awaittask
140+
self.assertFalse(task.cancelled())
141+
self.assertTrue(task.done())
142+
143+
asyncdeftest_nested_timeouts_concurrent(self):
144+
withself.assertRaises(TimeoutError):
145+
asyncwithasyncio.timeout(0.002):
146+
withself.assertRaises(TimeoutError):
147+
asyncwithasyncio.timeout(0.1):
148+
# Pretend we crunch some numbers.
149+
time.sleep(0.01)
150+
awaitasyncio.sleep(1)
151+
152+
asyncdeftest_nested_timeouts_loop_busy(self):
153+
# After the inner timeout is an expensive operation which should
154+
# be stopped by the outer timeout.
155+
loop=asyncio.get_running_loop()
156+
# Disable a message about long running task
157+
loop.slow_callback_duration=10
158+
t0=loop.time()
159+
withself.assertRaises(TimeoutError):
160+
asyncwithasyncio.timeout(0.1): # (1)
161+
withself.assertRaises(TimeoutError):
162+
asyncwithasyncio.timeout(0.01): # (2)
163+
# Pretend the loop is busy for a while.
164+
time.sleep(0.1)
165+
awaitasyncio.sleep(1)
166+
# TimeoutError was cought by (2)
167+
awaitasyncio.sleep(10) # This sleep should be interrupted by (1)
168+
t1=loop.time()
169+
self.assertTrue(t0<=t1<=t0+1)
170+
171+
asyncdeftest_reschedule(self):
172+
loop=asyncio.get_running_loop()
173+
fut=loop.create_future()
174+
deadline1=loop.time() +10
175+
deadline2=deadline1+20
176+
177+
asyncdeff():
178+
asyncwithasyncio.timeout_at(deadline1) ascm:
179+
fut.set_result(cm)
180+
awaitasyncio.sleep(50)
181+
182+
task=asyncio.create_task(f())
183+
cm=awaitfut
184+
185+
self.assertEqual(cm.when(), deadline1)
186+
cm.reschedule(deadline2)
187+
self.assertEqual(cm.when(), deadline2)
188+
cm.reschedule(None)
189+
self.assertIsNone(cm.when())
190+
191+
task.cancel()
192+
193+
withself.assertRaises(asyncio.CancelledError):
194+
awaittask
195+
self.assertFalse(cm.expired())
196+
197+
asyncdeftest_repr_active(self):
198+
asyncwithasyncio.timeout(10) ascm:
199+
self.assertRegex(repr(cm), r"<Timeout \[active\] when=\d+\.\d*>")
200+
201+
asyncdeftest_repr_expired(self):
202+
withself.assertRaises(TimeoutError):
203+
asyncwithasyncio.timeout(0.01) ascm:
204+
awaitasyncio.sleep(10)
205+
self.assertEqual(repr(cm), "<Timeout [expired]>")
206+
207+
asyncdeftest_repr_finished(self):
208+
asyncwithasyncio.timeout(10) ascm:
209+
awaitasyncio.sleep(0)
210+
211+
self.assertEqual(repr(cm), "<Timeout [finished]>")
212+
213+
asyncdeftest_repr_disabled(self):
214+
asyncwithasyncio.timeout(None) ascm:
215+
self.assertEqual(repr(cm), r"<Timeout [active] when=None>")
216+
217+
asyncdeftest_nested_timeout_in_finally(self):
218+
withself.assertRaises(TimeoutError):
219+
asyncwithasyncio.timeout(0.01):
220+
try:
221+
awaitasyncio.sleep(1)
222+
finally:
223+
withself.assertRaises(TimeoutError):
224+
asyncwithasyncio.timeout(0.01):
225+
awaitasyncio.sleep(10)
226+
227+
228+
if__name__=='__main__':
229+
unittest.main()

0 commit comments

Comments
(0)