stats: use try/catch for contextmanager

my/core/__main__.py

@@ -219,7 +219,7 @@ def modules_check(*, verbose: bool, list_all: bool, quick: bool, for_modules: Li
 
     import contextlib
 
-    from .common import with_quick_stats
+    from .common import quick_stats
     from .util import get_stats, HPIModule
     from .stats import guess_stats
     from .error import warn_my_config_import_error
@@ -263,7 +263,7 @@ def modules_check(*, verbose: bool, list_all: bool, quick: bool, for_modules: Li
             # todo point to a readme on the module structure or something?
             continue
 
-        quick_context = with_quick_stats() if quick else contextlib.nullcontext()
+        quick_context = quick_stats() if quick else contextlib.nullcontext()
 
         try:
             kwargs = {}

my/core/common.py

@@ -427,7 +427,7 @@ def warn_if_empty(f):
 
 
 # global state that turns on/off quick stats
-# can use the 'with_quick_stats' contextmanager
+# can use the 'quick_stats' contextmanager
 # to enable/disable this in cli so that module 'stats'
 # functions don't have to implement custom 'quick' logic
 QUICK_STATS = False
@@ -437,11 +437,14 @@ QUICK_STATS = False
 # elsewhere -- can use this decorator instead of editing
 # the global state directly
 @contextmanager
-def with_quick_stats():
+def quick_stats():
     global QUICK_STATS
-    prev, QUICK_STATS = QUICK_STATS, True
+    prev = QUICK_STATS
-    yield
+    try:
-    QUICK_STATS = prev
+        QUICK_STATS = True
+        yield
+    finally:
+        QUICK_STATS = prev
 
 
 C = TypeVar('C')