Add a profile decorator.
David Blume

David Blume commited on 2021-01-09 17:52:58
Showing 5 changed files, with 64 additions and 14 deletions.

... ...
@@ -1 +0,0 @@
1
-__all__ = ["counter", ]
... ...
@@ -0,0 +1,47 @@
1
+#!/usr/bin/env python3
2
+"""
3
+A profiler decorator.
4
+
5
+Author: Giampaolo Rodola' <g.rodola [AT] gmail [DOT] com>
6
+License: MIT
7
+"""
8
+import cProfile
9
+import tempfile
10
+import pstats
11
+from typing import Union, Optional, Tuple, List, Callable
12
+
13
+
14
+def profile(sort: Union[Tuple, List, str]='cumulative',
15
+            lines_to_print: Union[int, float]=20,
16
+            strip_dirs: bool=True,
17
+            filename: Optional[str]=None) -> Callable:
18
+    """A decorator which profiles a callable.
19
+    Output goes to stdout if filename is None, otherwise to filename."""
20
+    def outer(fun: Callable) -> Callable:
21
+        def inner(*args, **kwargs):
22
+            prof = cProfile.Profile()
23
+            ret = prof.runcall(fun, *args, **kwargs)
24
+            if filename is None:
25
+                stats = pstats.Stats(prof)
26
+            else:
27
+                statsfile = open(filename, 'a')
28
+                stats = pstats.Stats(prof, stream=statsfile)
29
+            if strip_dirs:
30
+                stats.strip_dirs()
31
+            if isinstance(sort, (tuple, list)):
32
+                stats.sort_stats(*sort)
33
+            else:
34
+                stats.sort_stats(sort)
35
+            stats.print_stats(lines_to_print)
36
+            if filename is not None:
37
+                statsfile.close()
38
+
39
+            return ret
40
+        return inner
41
+
42
+    # in case this is defined as "@profile" instead of "@profile()"
43
+    if hasattr(sort, '__call__'):
44
+        fun = sort
45
+        sort = 'cumulative'
46
+        outer = outer(fun)
47
+    return outer
... ...
@@ -0,0 +1,14 @@
1
+import time
2
+import functools
3
+from typing import Callable
4
+
5
+
6
+def timeit(f: Callable) -> Callable:
7
+    """Decorator that prints the duration of the decorated function."""
8
+    @functools.wraps(f)
9
+    def wrapper(*args, **kwargs):
10
+        s = time.time()
11
+        r = f(*args, **kwargs)
12
+        print(f'{f.__name__} took {time.time() - s:1.3f}s.')
13
+        return r
14
+    return wrapper
... ...
@@ -1 +0,0 @@
1
-__all__ = ["sitesize", ]
... ...
@@ -8,17 +8,7 @@ from argparse import ArgumentParser
8 8
 import counter.counter
9 9
 import sitesize.sitesize
10 10
 from typing import Optional, Callable
11
-
12
-
13
-def timeit(f: Callable) -> Callable:
14
-    """Decorator that prints the duration of the decorated function."""
15
-    @functools.wraps(f)
16
-    def wrapper(*args, **kwargs):
17
-        s = time.time()
18
-        r = f(*args, **kwargs)
19
-        print(f'{f.__name__} took {time.time() - s:1.3f}s.')
20
-        return r
21
-    return wrapper
11
+from decorators import timeit, profile
22 12
 
23 13
 
24 14
 v_print:Callable
... ...
@@ -50,7 +40,8 @@ class Coffee:
50 40
             raise ValueError("price must be a non-negative float.")
51 41
 
52 42
 
53
-@timeit
43
+@timeit.timeit
44
+@profile.profile  # may invoke with parameters, too
54 45
 def main(debug: bool) -> None:
55 46
     script_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
56 47
     print(f'{sys.argv[0]} is in {script_dir}.')
57 48