core: more type annotations

my/core/__init__.py

@@ -3,7 +3,7 @@ from .common import PathIsh, Paths, Json
 from .common import get_files
 from .common import LazyLogger
 from .common import warn_if_empty
-from .common import stat
+from .common import stat, Stats
 
 from .cfg import make_config
 from .util import __NOT_HPI_MODULE__

my/core/common.py

@@ -362,8 +362,9 @@ QUICK_STATS = False
 
 
 C = TypeVar('C')
+Stats = Dict[str, Any]
 # todo not sure about return type...
-def stat(func: Callable[[], Iterable[C]]) -> Dict[str, Any]:
+def stat(func: Callable[[], Iterable[C]]) -> Stats:
     from more_itertools import ilen, take, first
 
     # todo not sure if there is something in more_itertools to compute this?

my/core/pandas.py

@@ -2,18 +2,24 @@
 Various pandas helpers and convenience functions
 '''
 # todo not sure if belongs to 'core'. It's certainly 'more' core than actual modules, but still not essential
-from typing import Optional
-import warnings
+# NOTE: this file is meant to be importable without Pandas installed
+from typing import Optional, TYPE_CHECKING, Any
+from . import warnings
 
 
-# FIXME need to make sure check_dataframe decorator can be used without actually importing pandas
-# so need to move this import drom top level
-import pandas as pd # type: ignore
-
-# todo special warning type?
+if TYPE_CHECKING:
+    # this is kinda pointless at the moment, but handy to annotate DF returning methods now
+    # later will be unignored when they implement type annotations
+    import pandas as pd # type: ignore
+    # DataFrameT = pd.DataFrame
+    DataFrameT = Any
+else:
+    # in runtime, make it defensive so it works without pandas
+    DataFrameT = Any
 
 
 def check_dateish(s) -> Optional[str]:
+    import pandas as pd # type: ignore
     ctype = s.dtype
     if str(ctype).startswith('datetime64'):
         return None
@@ -26,18 +32,22 @@ def check_dateish(s) -> Optional[str]:
     return None
 
 
-def check_dataframe(f):
+from typing import Any, Callable, TypeVar
+FuncT = TypeVar('FuncT', bound=Callable[..., DataFrameT])
+
+def check_dataframe(f: FuncT) -> FuncT:
     from functools import wraps
     @wraps(f)
-    def wrapper(*args, **kwargs) -> pd.DataFrame:
+    def wrapper(*args, **kwargs) -> DataFrameT:
         df = f(*args, **kwargs)
         # todo make super defensive?
         # TODO check index as well?
         for col, data in df.iteritems():
             res = check_dateish(data)
             if res is not None:
-                warnings.warn(f"{f.__name__}, column '{col}': {res}")
+                warnings.low(f"{f.__name__}, column '{col}': {res}")
         return df
-    return wrapper
+    # https://github.com/python/mypy/issues/1927
+    return wrapper # type: ignore[return-value]
 
 # todo doctor: could have a suggesion to wrap dataframes with it?? discover by return type?

my/core/types.py

@@ -1,9 +0,0 @@
-import typing
-
-if typing.TYPE_CHECKING:
-    from typing import Any
-    # todo would be nice to use some real stubs..
-    DataFrameT = Any
-else:
-    import pandas # type: ignore
-    DataFrameT = pandas.DataFrame

my/core/warnings.py

@@ -35,6 +35,11 @@ def _warn(message: str, *args, color=None, **kwargs) -> None:
     warnings.warn(_colorize(message, color=color), *args, **kwargs)
 
 
+def low(message: str, *args, **kwargs) -> None:
+    kwargs['color'] = 'grey'
+    _warn(message, *args, **kwargs)
+
+
 def medium(message: str, *args, **kwargs) -> None:
     kwargs['color'] = 'yellow'
     _warn(message, *args, **kwargs)