doctor: better quick option propogation for stats (#239)

doctor: better quick option propogation for stats

* use contextmanager for quick stats instead of editing global state
  directly
* send quick to lots of stat related functions, so they
could possibly be used without doctor, if someone wanted to
* if a stats function has a 'quick' kwarg, send the value
there as well
* add an option to sort locations in my.time.tz.via_location

my/core/__main__.py

@@ -215,11 +215,11 @@ def modules_check(*, verbose: bool, list_all: bool, quick: bool, for_modules: Li
         verbose = True
     vw = '' if verbose else '; pass --verbose to print more information'
 
-    from . import common
-    common.QUICK_STATS = quick # dirty, but hopefully OK for cli
-
     tabulate_warnings()
 
+    import contextlib
+
+    from .common import quick_stats
     from .util import get_stats, HPIModule
     from .stats import guess_stats
     from .error import warn_my_config_import_error
@@ -256,15 +256,21 @@ def modules_check(*, verbose: bool, list_all: bool, quick: bool, for_modules: Li
         stats = get_stats(m)
         if stats is None:
             # then try guessing.. not sure if should log somehow?
-            stats = guess_stats(m)
+            stats = guess_stats(m, quick=quick)
 
         if stats is None:
             eprint("       - no 'stats' function, can't check the data")
             # todo point to a readme on the module structure or something?
             continue
 
+        quick_context = quick_stats() if quick else contextlib.nullcontext()
+
         try:
-            res = stats()
+            kwargs = {}
+            if callable(stats) and 'quick' in inspect.signature(stats).parameters:
+                kwargs['quick'] = quick
+            with quick_context:
+                res = stats(**kwargs)
             assert res is not None, 'stats() returned None'
         except Exception as ee:
             warning(f'     - {click.style("stats:", fg="red")}                      computing failed{vw}')

my/core/common.py

@@ -2,6 +2,7 @@ from glob import glob as do_glob
 from pathlib import Path
 from datetime import datetime
 import functools
+from contextlib import contextmanager
 import types
 from typing import Union, Callable, Dict, Iterable, TypeVar, Sequence, List, Optional, Any, cast, Tuple, TYPE_CHECKING
 import warnings
@@ -425,16 +426,32 @@ def warn_if_empty(f):
     return wrapped # type: ignore
 
 
-# hacky hook to speed up for 'hpi doctor'
-# todo think about something better
+# global state that turns on/off quick stats
+# 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
 
 
+# incase user wants to use the stats functions/quick option
+# elsewhere -- can use this decorator instead of editing
+# the global state directly
+@contextmanager
+def quick_stats():
+    global QUICK_STATS
+    prev = QUICK_STATS
+    try:
+        QUICK_STATS = True
+        yield
+    finally:
+        QUICK_STATS = prev
+
+
 C = TypeVar('C')
 Stats = Dict[str, Any]
 StatsFun = Callable[[], Stats]
 # todo not sure about return type...
-def stat(func: Union[Callable[[], Iterable[C]], Iterable[C]]) -> Stats:
+def stat(func: Union[Callable[[], Iterable[C]], Iterable[C]], quick: bool=False) -> Stats:
     if callable(func):
         fr = func()
         fname = func.__name__
@@ -451,13 +468,13 @@ def stat(func: Union[Callable[[], Iterable[C]], Iterable[C]]) -> Stats:
             rows=len(df),
         )
     else:
-        res = _stat_iterable(fr)
+        res = _stat_iterable(fr, quick=quick)
     return {
         fname: res,
     }
 
 
-def _stat_iterable(it: Iterable[C]) -> Any:
+def _stat_iterable(it: Iterable[C], quick: bool=False) -> Any:
     from more_itertools import ilen, take, first
 
     # todo not sure if there is something in more_itertools to compute this?
@@ -476,7 +493,7 @@ def _stat_iterable(it: Iterable[C]) -> Any:
 
     eit = funcit()
     count: Any
-    if QUICK_STATS:
+    if quick or QUICK_STATS:
         initial = take(100, eit)
         count = len(initial)
         if first(eit, None) is not None: # todo can actually be none...

my/core/stats.py

@@ -13,12 +13,12 @@ from .common import StatsFun, Stats, stat
 
 # TODO maybe could be enough to annotate OUTPUTS or something like that?
 # then stats could just use them as hints?
-def guess_stats(module_name: str) -> Optional[StatsFun]:
+def guess_stats(module_name: str, quick: bool=False) -> Optional[StatsFun]:
     providers = guess_data_providers(module_name)
     if len(providers) == 0:
         return None
     def auto_stats() -> Stats:
-        return {k: stat(v) for k, v in providers.items()}
+        return {k: stat(v, quick=quick) for k, v in providers.items()}
     return auto_stats
 
 

my/core/util.py

@@ -20,7 +20,7 @@ def get_stats(module: str) -> Optional[StatsFun]:
     # todo detect via ast?
     try:
         mod = import_module(module)
-    except Exception as e:
+    except Exception:
         return None
 
     return getattr(mod, 'stats', None)