From a57a06b1cc74eec75dcc9d44f3192c263eb76b2c Mon Sep 17 00:00:00 2001
From: jesopo
Date: Fri, 30 Aug 2019 18:13:12 +0100
Subject: switch utils.deadline to a context manager

---
 src/utils/__init__.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

(limited to 'src/utils')

diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index 35ac23e5..719065c0 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,4 +1,5 @@
-import datetime, decimal, enum, io, ipaddress, re, signal, threading, typing
+import contextlib, datetime, decimal, enum, io, ipaddress, re, signal
+import threading, typing
 from src.utils import cli, consts, irc, http, parse, security
 
 class Direction(enum.Enum):
@@ -348,11 +349,13 @@ class DeadlineExceededException(Exception):
 def _raise_deadline():
     raise DeadlineExceededException()
 
-def deadline(f: typing.Callable[[], typing.Any], seconds: int=10) -> typing.Any:
-    signal.signal(signal.SIGALRM, lambda _1, _2: _raise_deadline())
+@contextlib.contextmanager
+def deadline(seconds: int=10):
+    old_handler = signal.signal(signal.SIGALRM,
+        lambda _1, _2: _raise_deadline())
     signal.alarm(seconds)
 
     try:
-        return f()
+        yield
     finally:
-        signal.signal(signal.SIGALRM, signal.SIG_IGN)
+        signal.signal(signal.SIGALRM, old_handler)
-- 
cgit v1.3.1-10-gc9f91