Wrap timing in a context manager

asyncpg/connection.py

@@ -339,10 +339,9 @@ async def execute(self, query: str, *args, timeout: float=None) -> str:
         self._check_open()
 
         if not args:
-            start = time.monotonic()
-            result = await self._protocol.query(query, timeout)
-            elapsed = time.monotonic() - start
-            self._log_query(query, args, timeout, elapsed)
+            with utils.timer() as t:
+                result = await self._protocol.query(query, timeout)
+            self._log_query(query, args, timeout, t.elapsed)
             return result
 
         _, status, _ = await self._execute(
@@ -1724,27 +1723,25 @@ async def __execute(
         executor = lambda stmt, timeout: self._protocol.bind_execute(
             stmt, args, '', limit, return_status, timeout)
         timeout = self._protocol._get_timeout(timeout)
-        start = time.monotonic()
-        result, stmt = await self._do_execute(
-            query,
-            executor,
-            timeout,
-            record_class=record_class,
-            ignore_custom_codec=ignore_custom_codec,
-        )
-        elapsed = time.monotonic() - start
-        self._log_query(query, args, timeout, elapsed)
+        with utils.timer() as t:
+            result, stmt = await self._do_execute(
+                query,
+                executor,
+                timeout,
+                record_class=record_class,
+                ignore_custom_codec=ignore_custom_codec,
+            )
+        self._log_query(query, args, timeout, t.elapsed)
         return result, stmt
 
     async def _executemany(self, query, args, timeout):
         executor = lambda stmt, timeout: self._protocol.bind_execute_many(
             stmt, args, '', timeout)
         timeout = self._protocol._get_timeout(timeout)
         with self._stmt_exclusive_section:
-            start = time.monotonic()
-            result, _ = await self._do_execute(query, executor, timeout)
-            elapsed = time.monotonic() - start
-        self._log_query(query, args, timeout, elapsed)
+            with utils.timer() as t:
+                result, _ = await self._do_execute(query, executor, timeout)
+        self._log_query(query, args, timeout, t.elapsed)
         return result
 
     async def _do_execute(

asyncpg/utils.py

@@ -6,6 +6,7 @@
 
 
 import re
+import time
 
 
 def _quote_ident(ident):
@@ -43,3 +44,28 @@ async def _mogrify(conn, query, args):
     # Finally, replace $n references with text values.
     return re.sub(
         r'\$(\d+)\b', lambda m: textified[int(m.group(1)) - 1], query)
+
+
+class timer:
+    __slots__ = ('start', 'elapsed')
+
+    def __init__(self):
+        self.start = time.monotonic()
+        self.elapsed = None
+
+    @property
+    def current(self):
+        return time.monotonic() - self.start
+
+    def restart(self):
+        self.start = time.monotonic()
+
+    def stop(self):
+        self.elapsed = self.current
+
+    def __enter__(self):
+        self.restart()
+        return self
+
+    def __exit__(self, *exc):
+        self.stop()