Enables use of SubTracker as a function decorator

src/gpu_tracker/__init__.py

@@ -11,3 +11,4 @@
 
 from .tracker import Tracker
 from .sub_tracker import SubTracker
+from .sub_tracker import sub_track

src/gpu_tracker/sub_tracker.py

@@ -2,22 +2,25 @@
 import inspect
 import os
 import time
+import functools
 from ._helper_classes import _TrackingFile, _SubTrackerLog
 
 
 class SubTracker:
     """
     Context manager that logs to a file for the purposes of sub tracking a code block using the timestamps at which the codeblock begins and ends.
     Entering the context manager marks the beginning of the code block and exiting the context manager marks the end of the code block.
-    At the beginning of the codeblock, the ``SubTracker`` logs a row to a tablular file (".csv" or ".sqlite") that includes the timestamp along with a name for the code block and an indication of whether it is the start or end of the code bock.
+    At the beginning of the codeblock, the ``SubTracker`` logs a row to a tabular file (".csv" or ".sqlite") that includes the timestamp along with a name for the code block and an indication of whether it is the start or end of the code bock.
     This resulting file can be used alongside a tracking file created by a ``Tracker`` object for more granular analysis of specific code blocks.
 
     :ivar str code_block_name: The name of the code block being sub-tracked.
     :ivar str sub_tracking_file: The path to the file where the sub-tracking info is logged.
     """
-    def __init__(self, code_block_name: str | None = None, sub_tracking_file: str | None = None):
+    def __init__(
+            self, code_block_name: str | None = None, code_block_attribute: str | None = None, sub_tracking_file: str | None = None):
         """
-        :param code_block_name: The name of the code block within a ``Tracker`` context that is being sub-tracked. Defaults to the file path and line number where the SubTracker context is started.
+        :param code_block_name: The name of the code block within a ``Tracker`` context that is being sub-tracked. Defaults to the file path followed by a colon followed by the ``code_block_attribute``.
+        :param code_block_attribute: Only used if ``code_block_name`` is ``None``. Defaults to the line number where the SubTracker context is started.
         :param sub_tracking_file: The path to the file to log the time stamps of the code block being sub-tracked Defaults to the ID of the process where the SubTracker context is created and in CSV format.
         """
         if code_block_name is not None:
@@ -26,8 +29,8 @@ def __init__(self, code_block_name: str | None = None, sub_tracking_file: str |
             stack = inspect.stack()
             caller_frame = stack[1]
             file_path = os.path.abspath(caller_frame.filename)
-            line_number = caller_frame.lineno
-            self.code_block_name = f'{file_path}:{line_number}'
+            code_block_attribute = caller_frame.lineno if code_block_attribute is None else code_block_attribute
+            self.code_block_name = f'{file_path}:{code_block_attribute}'
         if sub_tracking_file is None:
             sub_tracking_file = f'{os.getpid()}.csv'
         self.sub_tracking_file = sub_tracking_file
@@ -44,3 +47,32 @@ def __enter__(self):
 
     def __exit__(self, *_):
         self._log(_SubTrackerLog.CodeBlockPosition.END)
+
+
+def sub_track(code_block_name: str | None = None, code_block_attribute: str | None = None, sub_tracking_file: str | None = None):
+    """
+    Decorator for sub tracking calls to a specified function.
+
+    :param code_block_name: The name of the code block within a ``Tracker`` context that is being sub-tracked. Defaults to the file path followed by a colon followed by the ``code_block_attribute``.
+    :param code_block_attribute: Only used if ``code_block_name`` is ``None``. Defaults to the name of the decorated function i.e. the function being sub-tracked.
+    :param sub_tracking_file: The path to the file to log the time stamps of the code block being sub-tracked. Defaults to the ID of the process where the SubTracker context is created and in CSV format.
+    """
+    def decorator(func):
+        nonlocal code_block_name, code_block_attribute, sub_tracking_file
+        if code_block_name is None:
+            stack = inspect.stack()
+            caller_frame = stack[1]
+            file_path = os.path.abspath(caller_frame.filename)
+            code_block_attribute = func.__name__ if code_block_attribute is None else code_block_attribute
+            code_block_name = f'{file_path}:{code_block_attribute}'
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            nonlocal sub_tracking_file
+            with SubTracker(
+                    code_block_name=code_block_name, code_block_attribute=code_block_attribute, sub_tracking_file=sub_tracking_file
+            ):
+                return_value = func(*args, **kwargs)
+            return return_value
+        return wrapper
+    return decorator

tests/data/decorated-function-other-file.csv

@@ -0,0 +1,3 @@
+position,timestamp
+START,12
+END,13

tests/data/None_sub-tracking-file.csv.csv → tests/data/decorated-function.csv

@@ -9,3 +9,5 @@ START,6
 END,7
 START,8
 END,9
+START,10
+END,11

tests/test_sub_tracker.py

@@ -34,8 +34,53 @@ def code_block_name_test(val: str):
             assert val.endswith(default_code_block_end)
         else:
             assert val == code_block_name
-    expected_tracking_file = f'tests/data/{code_block_name}_{sub_tracking_file}.csv'
     utils.test_tracking_file(
-        actual_tracking_file=sub_tracker.sub_tracking_file, expected_tracking_file=expected_tracking_file,
+        actual_tracking_file=sub_tracker.sub_tracking_file, expected_tracking_file='tests/data/sub-tracker.csv',
         excluded_col='code_block_name', excluded_col_test=code_block_name_test
     )
+
+
+@pt.fixture(name='code_block_attribute', params=['my-attribute', None])
+def get_code_block_attribute(request):
+    yield request.param
+
+
+def test_decorator(mocker, code_block_name: str | None, code_block_attribute: str | None):
+    @gput.sub_track(code_block_name=code_block_name, code_block_attribute=code_block_attribute)
+    def decorated_function(arg1: int, arg2: int, kwarg1: int = 1, kwarg2: int = 2) -> int:
+        return arg1 + arg2 - (kwarg1 + kwarg2)
+    getpid_mock = mocker.patch('gpu_tracker.sub_tracker.os.getpid', return_value=1234)
+    n_iterations = 3
+    time_mock = mocker.patch('gpu_tracker.sub_tracker.time', time=mocker.MagicMock(side_effect=range(n_iterations * 2 * 2 + 2)))
+    for _ in range(n_iterations):
+        return_val = decorated_function(2, 5)
+        assert return_val == 4
+        return_val = decorated_function(3, 2, kwarg1=2, kwarg2=1)
+        assert return_val == 2
+    assert len(getpid_mock.call_args_list) == n_iterations * 2
+    assert len(time_mock.time.call_args_list) == n_iterations * 2 * 2
+
+    def code_block_name_test(val):
+        if code_block_name is None:
+            if code_block_attribute is None:
+                assert val.endswith('test_sub_tracker.py:decorated_function')
+            else:
+                assert val.endswith('test_sub_tracker.py:my-attribute')
+        else:
+            assert val == code_block_name
+    utils.test_tracking_file(
+        actual_tracking_file='1234.csv', expected_tracking_file=f'tests/data/decorated-function.csv',
+        excluded_col='code_block_name', excluded_col_test=code_block_name_test
+    )
+    if code_block_name is None and code_block_attribute is None:
+        return_val = utils.function_in_other_file(1, 2, 3, kw1=4, kw2=5)
+        assert return_val == ((1, 2, 3), {'kw1': 4, 'kw2': 5})
+        assert len(getpid_mock.call_args_list) == n_iterations * 2 + 1
+        assert len(time_mock.time.call_args_list) == n_iterations * 2 * 2 + 2
+
+        def code_block_name_test(val):
+            assert val.endswith('utils.py:function_in_other_file')
+        utils.test_tracking_file(
+            actual_tracking_file='1234.csv', expected_tracking_file='tests/data/decorated-function-other-file.csv',
+            excluded_col='code_block_name', excluded_col_test=code_block_name_test
+        )

tests/utils.py

@@ -3,6 +3,7 @@
 import os
 # noinspection PyProtectedMember
 from gpu_tracker._helper_classes import _SQLiteTrackingFile
+import gpu_tracker as gput
 
 
 def assert_args_list(mock, expected_args_list: list[tuple | dict], use_kwargs: bool = False):
@@ -23,3 +24,8 @@ def test_tracking_file(
     expected_tracking_log = pd.read_csv(expected_tracking_file)
     pd.testing.assert_frame_equal(expected_tracking_log, actual_tracking_log, atol=1e-10, rtol=1e-10)
     os.remove(actual_tracking_file)
+
+
+@gput.sub_track()
+def function_in_other_file(*args, **kwargs):
+    return args, kwargs