add: performance in toad.utils

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Add
+- Added `performance` in `toad.utils` for test code performance
+
 ## [0.1.2] - 2023-04-09
 
 ### Add

toad/__init__.py

@@ -12,7 +12,7 @@
 from .stats import quality, IV, VIF, WOE, entropy, entropy_cond, gini, gini_cond
 from .selection import select
 from .scorecard import ScoreCard
-from .utils import Progress
+from .utils import Progress, performance
 from .version import __version__
 
 VERSION = __version__

toad/utils/decorator.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+from time import time
 from .func import save_json, read_json
 from functools import wraps, WRAPPER_ASSIGNMENTS
 
@@ -230,3 +231,62 @@ def partial_func(x):
 
         return grad, hess
 
+
+class performance(Decorator):
+    """decorator for analysis code performance
+
+    Args:
+        loop (int): loop times, default `1`
+    
+    Examples:
+    >>> @performance(loop = 100)
+    >>> def func():
+    >>>     ... # code
+    >>>     return res
+    >>>
+    >>> func()
+    >>> 
+    >>> # or use `performance` in `with` statement
+    >>> with performance():
+    >>>     ... # code
+    """
+    loop = 1
+
+    def wrapper(self, *args, **kwargs):
+        costs = []
+        for _ in range(self.loop):
+            start = time()
+            res = self.call(*args, **kwargs)
+            end = time()
+            costs.append(end - start)
+
+        self.analysis(costs)
+        return res
+
+
+    def analysis(self, costs):
+        import numpy as np
+
+        print('total cost: {:.5f}s'.format(np.sum(costs)))
+        print("-"*40)
+        data = {
+            "Mean": np.mean(costs),
+            "Min": np.min(costs),
+            "Max": np.max(costs),
+            "90%": np.percentile(costs, 90),
+            "95%": np.percentile(costs, 95),
+            "99%": np.percentile(costs, 99),
+        }
+        HEADER = "{:>8}"*len(data)
+        BODY = "{:>7.3f}s"*len(data)
+        print(HEADER.format(*data.keys()))
+        print(BODY.format(*data.values()))
+
+
+    def __enter__(self):
+        self.start = time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end = time()
+        self.analysis([self.end - self.start])

toad/utils/decorator_test.py

@@ -2,7 +2,12 @@
 import numpy as np
 import pandas as pd
 
-from .decorator import Decorator, frame_exclude, xgb_loss
+from .decorator import (
+    Decorator,
+    frame_exclude,
+    xgb_loss,
+    performance,
+)
 
 np.random.seed(1)
 
@@ -53,3 +58,25 @@ def loss(x, y):
 
     assert grad == pytest.approx(-3.0)
     assert hess == pytest.approx(0.0)
+
+
+def test_performance():
+    @performance(loop = 10)
+    def func(x):
+        from time import sleep
+        sleep(0.01)
+        return x**x
+
+    assert func(2) == 4
+
+
+def test_performance_with_clause():
+    def func(x):
+        from time import sleep
+        sleep(0.01)
+        return x**x
+
+    with performance():
+        res = func(2)
+
+    assert res == 4