Wrap timing in a context manager

Author: dcwatson

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()