diff --git a/docs/index.rst b/docs/index.rst index cbeaa40..e71195d 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/timed.md .. toctree:: diff --git a/docs/md/timed.md b/docs/md/timed.md new file mode 100644 index 0000000..d4de7ab --- /dev/null +++ b/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. \ No newline at end of file diff --git a/modflow_devtools/misc.py b/modflow_devtools/misc.py index 7dd9210..9646aee 100644 --- a/modflow_devtools/misc.py +++ b/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 diff --git a/modflow_devtools/test/test_misc.py b/modflow_devtools/test/test_misc.py index 79ea786..ff4f980 100644 --- a/modflow_devtools/test/test_misc.py +++ b/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)