core: more type annotations

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?