Revert "Revert "feat(timeit): add function timing decorator (#118)""

docs/index.rst

@@ -31,6 +31,7 @@ The `modflow-devtools` package provides a set of tools for developing and testin
    md/download.md
    md/ostags.md
    md/zip.md
+   md/timed.md
 
 
 .. toctree::

docs/md/timed.md

@@ -0,0 +1,21 @@
+# `timed`
+
+There is a `@timed` decorator function available in the `modflow_devtools.misc` module. Applying it to any function prints a (rough) benchmark to `stdout` when the function returns. For instance:
+
+```python
+from modflow_devtools.misc import timed
+
+@timed
+def sleep1():
+    sleep(0.001)
+
+sleep1() # prints e.g. "sleep1 took 1.26 ms"
+```
+
+It can also wrap a function directly:
+
+```python
+timed(sleep1)()
+```
+
+The [`timeit`](https://docs.python.org/3/library/timeit.html) built-in module is used internally, however the timed function is only called once, where by default, `timeit` averages multiple runs.

modflow_devtools/misc.py

@@ -3,12 +3,14 @@
 import sys
 import traceback
 from contextlib import contextmanager
+from functools import wraps
 from importlib import metadata
 from os import PathLike, chdir, environ, getcwd
 from os.path import basename, normpath
 from pathlib import Path
 from shutil import which
 from subprocess import PIPE, Popen
+from timeit import timeit
 from typing import List, Optional, Tuple
 from urllib import request
 
@@ -442,3 +444,44 @@ def try_metadata() -> bool:
     _has_pkg_cache[pkg] = found
 
     return _has_pkg_cache[pkg]
+
+
+def timed(f):
+    """
+    Decorator for estimating runtime of any function.
+    Prints estimated time to stdout, in milliseconds.
+
+    Parameters
+    ----------
+    f : function
+        Function to time.
+
+    Notes
+    -----
+    Adapted from https://stackoverflow.com/a/27737385/6514033.
+    Uses the built-in timeit module internally.
+
+    Returns
+    -------
+    function
+        The decorated function.
+    """
+
+    @wraps(f)
+    def _timed(*args, **kw):
+        res = None
+
+        def call():
+            nonlocal res
+            res = f(*args, **kw)
+
+        t = timeit(lambda: call(), number=1)
+        if "log_time" in kw:
+            name = kw.get("log_name", f.__name__.upper())
+            kw["log_time"][name] = int(t * 1000)
+        else:
+            print(f"{f.__name__} took {t * 1000:.2f} ms")
+
+        return res
+
+    return _timed

modflow_devtools/test/test_misc.py

@@ -1,8 +1,9 @@
 import os
+import re
 import shutil
 from os import environ
 from pathlib import Path
-from pprint import pprint
+from time import sleep
 from typing import List
 
 import pytest
@@ -12,9 +13,9 @@
     get_namefile_paths,
     get_packages,
     has_package,
-    has_pkg,
     set_dir,
     set_env,
+    timed,
 )
 
 
@@ -25,7 +26,7 @@ def test_set_dir(tmp_path):
     assert Path(os.getcwd()) != tmp_path
 
 
-def test_set_env(tmp_path):
+def test_set_env():
     # test adding a variable
     key = "TEST_ENV"
     val = "test"
@@ -283,3 +284,24 @@ def test_has_pkg(virtualenv):
             ).strip()
             == exp
         )
+
+
+def test_timed1(capfd):
+    def sleep1():
+        sleep(0.001)
+
+    timed(sleep1)()
+    cap = capfd.readouterr()
+    print(cap.out)
+    assert re.match(r"sleep1 took \d+\.\d+ ms", cap.out)
+
+
+def test_timed2(capfd):
+    @timed
+    def sleep1dec():
+        sleep(0.001)
+
+    sleep1dec()
+    cap = capfd.readouterr()
+    print(cap.out)
+    assert re.match(r"sleep1dec took \d+\.\d+ ms", cap.out)