From 506a238f6f31d827015a6c6f5ba1867ee55948a7 Mon Sep 17 00:00:00 2001
From: wpbonelli <wbonelli@ucar.edu>
Date: Tue, 12 Sep 2023 18:03:19 -0400
Subject: [PATCH] feat(timeit): add function timing decorator (#118)

---
 docs/index.rst                     |  1 +
 docs/md/timeit.md                  | 17 ++++++++++++++
 modflow_devtools/misc.py           | 37 ++++++++++++++++++++++++++++++
 modflow_devtools/test/test_misc.py | 24 +++++++++++++++++++
 4 files changed, 79 insertions(+)
 create mode 100644 docs/md/timeit.md

diff --git a/docs/index.rst b/docs/index.rst
index cbeaa40..393f81b 100644
--- a/docs/index.rst
+++ b/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/timeit.md
 
 
 .. toctree::
diff --git a/docs/md/timeit.md b/docs/md/timeit.md
new file mode 100644
index 0000000..c57ef85
--- /dev/null
+++ b/docs/md/timeit.md
@@ -0,0 +1,17 @@
+# `timeit`
+
+There is a `timeit` decorator function available in the `modflow_devtools.misc` module. Applying it to any function causes a (rough) runtime benchmark to be printed to `stdout` afterwards the function returns. For instance:
+
+```python
+@timeit
+def sleep1():
+    sleep(0.001)
+
+sleep1() # prints e.g. "sleep1 took 1.26 ms"
+```
+
+`timeit` can also directly wrap a function:
+
+```python
+timeit(sleep1)() # prints same as above
+```
\ No newline at end of file
diff --git a/modflow_devtools/misc.py b/modflow_devtools/misc.py
index 7dd9210..284a6ed 100644
--- a/modflow_devtools/misc.py
+++ b/modflow_devtools/misc.py
@@ -1,8 +1,10 @@
 import importlib
 import socket
 import sys
+import time
 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
@@ -442,3 +444,38 @@ def try_metadata() -> bool:
     _has_pkg_cache[pkg] = found
 
     return _has_pkg_cache[pkg]
+
+
+def timeit(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.
+
+    Returns
+    -------
+    function
+        The decorated function.
+    """
+
+    @wraps(f)
+    def timed(*args, **kw):
+        ts = time.time()
+        res = f(*args, **kw)
+        te = time.time()
+        if "log_time" in kw:
+            name = kw.get("log_name", f.__name__.upper())
+            kw["log_time"][name] = int((te - ts) * 1000)
+        else:
+            print(f"{f.__name__} took {(te - ts) * 1000:.2f} ms")
+        return res
+
+    return timed
diff --git a/modflow_devtools/test/test_misc.py b/modflow_devtools/test/test_misc.py
index 79ea786..b8ebd40 100644
--- a/modflow_devtools/test/test_misc.py
+++ b/modflow_devtools/test/test_misc.py
@@ -1,8 +1,10 @@
 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
@@ -15,6 +17,7 @@
     has_pkg,
     set_dir,
     set_env,
+    timeit,
 )
 
 
@@ -283,3 +286,24 @@ def test_has_pkg(virtualenv):
             ).strip()
             == exp
         )
+
+
+def test_timeit1(capfd):
+    def sleep1():
+        sleep(0.001)
+
+    timeit(sleep1)()
+    cap = capfd.readouterr()
+    print(cap.out)
+    assert re.match(r"sleep1 took \d+\.\d+ ms", cap.out)
+
+
+def test_timeit2(capfd):
+    @timeit
+    def sleep1dec():
+        sleep(0.001)
+
+    sleep1dec()
+    cap = capfd.readouterr()
+    print(cap.out)
+    assert re.match(r"sleep1dec took \d+\.\d+ ms", cap.out)