Use null context manager to manage collection of statistics or not.

pull/3284/head
derekpierre 2023-10-13 19:08:47 -04:00
parent 0048b76506
commit 3d9fa8263f
1 changed files with 4 additions and 9 deletions

View File

@ -1,6 +1,7 @@
import cProfile
import os
import pstats
from contextlib import nullcontext
from nucypher_core.ferveo import DkgPublicKey
@ -89,12 +90,7 @@ bob = Bob(
bob.start_learning_loop(now=True)
profiler = None
try:
if collect_stats:
profiler = cProfile.Profile()
profiler.enable()
with cProfile.Profile() if collect_stats else nullcontext() as profiler:
cleartext = bob.threshold_decrypt(
threshold_message_kit=threshold_message_kit,
)
@ -102,9 +98,8 @@ try:
cleartext = bytes(cleartext)
print(f"\nCleartext: {cleartext.decode()}")
assert message == cleartext
finally:
if profiler:
profiler.disable()
if collect_stats:
profiler_stats = pstats.Stats(profiler).sort_stats(pstats.SortKey.TIME)
print("\n------ Profile Stats -------")
profiler_stats.print_stats(10)