Adding decorators.py.

pism_ragis/decorators.py

@@ -0,0 +1,120 @@
+# Copyright (C) 2024 Andy Aschwanden
+#
+# This file is part of pism-ragis.
+#
+# PISM-RAGIS is free software; you can redistribute it and/or modify it under the
+# terms of the GNU General Public License as published by the Free Software
+# Foundation; either version 3 of the License, or (at your option) any later
+# version.
+#
+# PISM-RAGIS is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
+# details.
+#
+# You should have received a copy of the GNU General Public License
+# along with PISM; if not, write to the Free Software
+
+# pylint: disable=too-many-positional-arguments
+
+"""
+Module for handling logging
+"""
+
+import logging
+import time
+import tracemalloc
+from functools import wraps
+
+from pism_ragis.logger import get_logger
+
+logger: logging.Logger = get_logger(__name__)
+
+
+def timeit(func):
+    """
+    Decorator that logs the time a function takes to execute.
+
+    This decorator logs the start time, end time, and the elapsed time
+    for the execution of the decorated function.
+
+    Parameters
+    ----------
+    func : callable
+        The function to be decorated.
+
+    Returns
+    -------
+    callable
+        The wrapped function with added timing functionality.
+
+    Examples
+    --------
+    >>> @timeit
+    ... def example_function():
+    ...     time.sleep(1)
+    ...
+    >>> example_function()
+    INFO:__main__:Starting example_function
+    INFO:__main__:Finished example_function in 1.0001 seconds
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        logger.info("%s: Starting", func.__name__)
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        logger.info("%s: Finished in %2.2f seconds", func.__name__, elapsed_time)
+        return result
+
+    return wrapper
+
+
+def profileit(func):
+    """
+    Decorator that logs the time and memory usage of a function.
+
+    This decorator logs the start time, end time, elapsed time, current memory usage,
+    and peak memory usage for the execution of the decorated function.
+
+    Parameters
+    ----------
+    func : callable
+        The function to be decorated.
+
+    Returns
+    -------
+    callable
+        The wrapped function with added profiling functionality.
+
+    Examples
+    --------
+    >>> @profileit
+    ... def example_function():
+    ...     time.sleep(1)
+    ...
+    >>> example_function()
+    INFO:__main__:Starting example_function
+    INFO:__main__:example_function: Memory usage: 0.10 MB
+    INFO:__main__:example_function: Peak memory usage: 0.15 MB
+    INFO:__main__:example_function Finished in 1.00 seconds
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        tracemalloc.start()
+        start_time = time.time()
+        logger.info("Starting %s", func.__name__)
+        result = func(*args, **kwargs)
+        current, peak = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        logger.info("%s: Memory usage: %.2f MB", func.__name__, current / 1024**2)
+        logger.info("%s: Peak memory usage: %.2f MB", func.__name__, peak / 1024**2)
+        logger.info("%s Finished in %.2f seconds", func.__name__, elapsed_time)
+        return result
+
+    return wrapper