[ Python 自製工具 ] count_fc.py: Count modul/function calls

 Preface

在某個工作下, 剛好需要統計在某個複雜程式執行中, 到底執行了多少模組底下的函數. 因此就開發了這個工具:
- count_fc.py

view plaincopy to clipboardprint?

1. import os  
2. import atexit  
3. import inspect  
4. import json  
5. import threading  
6. import logging  
7. import coloredlogs  
8. from collections import defaultdict  
9.   
10. ################################  
11. # Constants  
12. ################################  
13. MODU_PATH = os.path.dirname(__file__) if os.path.dirname(__file__) else './'  
14. ''' Path of current module '''  
15.   
16. LOGGER_FORMAT = "[%(levelno)s|%(module)s|%(lineno)s] %(message)s"  
17. ''' Format of Logger '''  
18.   
19. LOGGER_LEVEL = 10  # CRITICAL=50; ERROR=40; WARNING=30; INFO=20; DEBUG=10  
20. ''' Message level of Logger '''  
21.   
22.   
23. ################################  
24. # Local variables  
25. ################################  
26. _glock = threading.Lock()  
27. ''' Global threading lock '''  
28. _unloaded_modu_set = set()  
29. ''' Used to keep already unloaded module to avoid duplicate unload '''  
30. _log = logging.getLogger(__name__)  
31. ''' Logger object '''  
32. _log.setLevel(LOGGER_LEVEL)  
33. _log.propagate = False  
34. coloredlogs.install(level=LOGGER_LEVEL, logger=_log, fmt=LOGGER_FORMAT)  
35.   
36.   
37. def func_count(f):  
38.   """Function decorator to accumulate the times of execution.  
39.   
40.   Args:  
41.     f: callable, function to count execution.  
42.   
43.   Returns:  
44.     The wrapped function.  
45.   """  
46.   def wrapped(*args, **kwargs):  
47.     wrapped.calls += 1  
48.     wrapped.name = f.__name__  
49.     wrapped.modu_name = f.__module__  
50.     wrapped.is_wrapped = True  
51.     return f(*args, **kwargs)  
52.   
53.   wrapped.calls = 0  
54.   return wrapped  
55.   
56.   
57. def unload(modu):  
58.   """The function to be called when a module is unloaded  
59.   
60.   Args:  
61.     modu: The module which is about to be unloaded.  
62.   """  
63.   with _glock:  
64.     if modu in _unloaded_modu_set:  
65.       return  
66.   
67.     _log.debug(f"Unload {modu}")  
68.     stat_pkfile_path = f'/tmp/{modu.__name__}.json'  
69.     flist = [a for a in dir(modu)]  
70.     fdict = defaultdict(int)  
71.     for a in flist:  
72.       fo = getattr(modu, a)  
73.       if hasattr(fo, 'is_wrapped') and fo.modu_name == modu.__name__:  
74.         _log.debug(f"\tCount {fo.name}")  
75.         fdict[fo.name] += 1  
76.   
77.     if fdict:  
78.       with open(stat_pkfile_path, 'w') as fw:  
79.         json.dump(fdict, fw)  
80.   
81.     _unloaded_modu_set.add(modu)  
82.   
83.   
84. def decorate_mod(modu):  
85.   """Decorate all the function(s) of input module  
86.   
87.   Args:  
88.     modu: This function will decorate all the function(s)  
89.       collected from this module with `func_count` and register  
90.       this module with a method which will be called when unload.  
91.   """  
92.   _log.debug(f'Decorate {modu}')  
93.   flist = [a for a in dir(modu)]  
94.   for a in flist:  
95.     fo = getattr(modu, a)  
96.     if inspect.isfunction(fo) and fo.__module__ == modu.__name__:  
97.       _log.debug(f"\tfo={fo}; {fo.__module__}, {modu.__name__}")  
98.       setattr(modu, fo.__name__, func_count(fo))  
99.   
100.     setattr(modu, 'unload', unload)  
101.     atexit.register(modu.unload, modu)  

Usage:
這邊以模組 xyz.py 與 haha.py 為例:
- haha.py

view plaincopy to clipboardprint?

1. from count_fc import decorate_mod  
2. from xyz import *  
3.   
4. def say_hi(name):  
5.   print(message(name))  
6.   
7. decorate_mod(sys.modules[__name__])  

- xyz.py

view plaincopy to clipboardprint?

1. import sys  
2. import random  
3. from count_fc import decorate_mod  
4.   
5.   
6.   
7. def message(name):  
8.   return random.choice(greeting_metods)(name)  
9.   
10.   
11. def greeting_with_aloha(name):  
12.   return f"Aloha, {name}"  
13.   
14.   
15. def greeting_with_hello(name):  
16.   return f"Hello, {name}"  
17.   
18.   
19. def greeting_with_hi(name):  
20.   return f"Hi, {name}"  
21.   
22. decorate_mod(sys.modules[__name__])  
23.   
24. greeting_metods = [  
25.   greeting_with_aloha,  
26.   greeting_with_hello,  
27.   greeting_with_hi  
28. ]  

接著是主程式:
- main.py

view plaincopy to clipboardprint?

1. #!/usr/bin/env python  
2. from haha import say_hi  
3.   
4. if __name__ == '__main__':  
5.   for name in ['John', 'Mary', 'Key', 'Peter', 'Jane', 'Bob']:  
6.     say_hi(name)  
7.   
8.   print(say_hi.calls)  

執行結果如下:

# ./main.py
[10|count_fc|92] Decorate <module 'xyz' from '/root/abc/xyz.py'>
[10|count_fc|97] fo=<function greeting_with_aloha at 0x7f2c6ceacd30>; xyz, xyz
[10|count_fc|97] fo=<function greeting_with_hello at 0x7f2c6ceacdc0>; xyz, xyz
[10|count_fc|97] fo=<function greeting_with_hi at 0x7f2c6ceace50>; xyz, xyz
[10|count_fc|97] fo=<function message at 0x7f2c6ceacca0>; xyz, xyz
[10|count_fc|92] Decorate <module 'haha' from '/root/abc/haha.py'>
[10|count_fc|97] fo=<function say_hi at 0x7f2c6cff7ee0>; haha, haha
Hello, John
Aloha, Mary
Aloha, Key
Hi, Peter
Aloha, Jane
Aloha, Bob
6
[10|count_fc|67] Unload <module 'haha' from '/root/abc/haha.py'>
[10|count_fc|74] Count say_hi
[10|count_fc|67] Unload <module 'xyz' from '/root/abc/xyz.py'>
[10|count_fc|74] Count greeting_with_aloha
[10|count_fc|74] Count greeting_with_hello
[10|count_fc|74] Count greeting_with_hi
[10|count_fc|74] Count message

接著可以去 /tmp/ 目錄撿收產生的每個模組的呼叫結果:

# ls -hl /tmp/*.json
-rw-r--r-- 1 root root 13 Jul 3 02:10 /tmp/haha.json
-rw-r--r-- 1 root root 36 Jul 3 02:10 /tmp/xyz.json

# cat /tmp/haha.json
{"say_hi": 1}

# cat /tmp/xyz.json
{"greeting_with_aloha": 4, "greeting_with_hello": 1, "greeting_with_hi": 1, "message": 6}