Missing return types for major Koalas objects (#1900)

Author: xinrong-meng

databricks/koalas/frame.py

@@ -43,6 +43,7 @@
     cast,
 )
 
+import matplotlib
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_list_like, is_dict_like, is_scalar
@@ -843,12 +844,12 @@ def add(self, other) -> "DataFrame":
     # create accessor for Koalas specific methods.
     koalas = CachedAccessor("koalas", KoalasFrameMethods)
 
-    def hist(self, bins=10, **kwds):
+    def hist(self, bins=10, **kwds) -> matplotlib.axes.Axes:
         return self.plot.hist(bins, **kwds)
 
     hist.__doc__ = KoalasPlotAccessor.hist.__doc__
 
-    def kde(self, bw_method=None, ind=None, **kwds):
+    def kde(self, bw_method=None, ind=None, **kwds) -> matplotlib.axes.Axes:
         return self.plot.kde(bw_method, ind, **kwds)
 
     kde.__doc__ = KoalasPlotAccessor.kde.__doc__
@@ -1432,7 +1433,7 @@ def items(self) -> Iterator:
         """This is an alias of ``iteritems``."""
         return self.iteritems()
 
-    def to_clipboard(self, excel=True, sep=None, **kwargs):
+    def to_clipboard(self, excel=True, sep=None, **kwargs) -> None:
         """
         Copy object to the system clipboard.
 
@@ -3806,7 +3807,7 @@ def nunique(
 
             return first_series(DataFrame(internal).transpose())
 
-    def round(self, decimals=0):
+    def round(self, decimals=0) -> "DataFrame":
         """
         Round a DataFrame to a variable number of decimal places.
 
@@ -8472,7 +8473,7 @@ def stack(self) -> Union["DataFrame", "ks.Series"]:
         else:
             return kdf
 
-    def unstack(self):
+    def unstack(self) -> Union["DataFrame", "ks.Series"]:
         """
         Pivot the (necessarily hierarchical) index labels.
 
@@ -9649,7 +9650,7 @@ def idxmin(self, axis=0) -> "ks.Series":
 
         return ks.from_pandas(kdf._to_internal_pandas().idxmin())  # type: ignore
 
-    def info(self, verbose=None, buf=None, max_cols=None, null_counts=None):
+    def info(self, verbose=None, buf=None, max_cols=None, null_counts=None) -> None:
         """
         Print a concise summary of a DataFrame.
 
@@ -9750,7 +9751,7 @@ def info(self, verbose=None, buf=None, max_cols=None, null_counts=None):
                 # hack to use pandas' info as is.
                 self._data = self
                 count_func = self.count
-                self.count = lambda: count_func().to_pandas()
+                self.count = lambda: count_func().to_pandas()  # type: ignore
                 return pd.DataFrame.info(
                     self,
                     verbose=verbose,
@@ -9761,7 +9762,7 @@ def info(self, verbose=None, buf=None, max_cols=None, null_counts=None):
                 )
             finally:
                 del self._data
-                self.count = count_func
+                self.count = count_func  # type: ignore
 
     # TODO: fix parameter 'axis' and 'numeric_only' to work same as pandas'
     def quantile(
@@ -10083,7 +10084,7 @@ class  max_speed
         else:
             return self.iloc[:, indices]  # type: ignore
 
-    def eval(self, expr, inplace=False) -> Optional["DataFrame"]:
+    def eval(self, expr, inplace=False) -> Optional[Union["DataFrame", "ks.Series"]]:
         """
         Evaluate a string describing operations on DataFrame columns.
 

databricks/koalas/generic.py

@@ -22,7 +22,7 @@
 from collections.abc import Iterable
 from distutils.version import LooseVersion
 from functools import reduce
-from typing import List, Optional, Tuple, Union, TYPE_CHECKING, cast
+from typing import Any, List, Optional, Tuple, Union, TYPE_CHECKING, cast
 import warnings
 
 import numpy as np  # noqa: F401
@@ -58,6 +58,10 @@ class Frame(object, metaclass=ABCMeta):
     The base class for both DataFrame and Series.
     """
 
+    @abstractmethod
+    def __getitem__(self, key):
+        pass
+
     @property
     @abstractmethod
     def _internal(self) -> InternalFrame:
@@ -405,7 +409,7 @@ def get_dtype_counts(self) -> pd.Series:
             dtypes = list(self.dtypes)
         return pd.Series(dict(Counter([d.name for d in dtypes])))
 
-    def pipe(self, func, *args, **kwargs):
+    def pipe(self, func, *args, **kwargs) -> Any:
         r"""
         Apply func(self, \*args, \*\*kwargs).
 
@@ -1683,7 +1687,7 @@ def bool(self) -> bool:
             raise TypeError("bool() expects DataFrame or Series; however, " "got [%s]" % (self,))
         return df.head(2)._to_internal_pandas().bool()
 
-    def first_valid_index(self):
+    def first_valid_index(self) -> Union[Any, Tuple[Any, ...]]:
         """
         Retrieves the index of the first valid value.
 
@@ -1778,7 +1782,7 @@ def first_valid_index(self):
 
         return first_valid_idx
 
-    def last_valid_index(self):
+    def last_valid_index(self) -> Union[Any, Tuple[Any, ...]]:
         """
         Return index for last non-NA/null value.
 
@@ -2024,7 +2028,7 @@ def expanding(self, min_periods=1) -> Expanding:
         """
         return Expanding(self, min_periods=min_periods)
 
-    def get(self, key, default=None):
+    def get(self, key, default=None) -> Any:
         """
         Get item from object for given key (DataFrame column, Panel slice,
         etc.). Returns default value if not found.

databricks/koalas/groupby.py

@@ -2657,10 +2657,10 @@ def _agg_columns(self):
     def _reduce_for_stat_function(self, sfun, only_numeric):
         return first_series(super()._reduce_for_stat_function(sfun, only_numeric))
 
-    def agg(self, *args, **kwargs):
+    def agg(self, *args, **kwargs) -> None:
         return MissingPandasLikeSeriesGroupBy.agg(self, *args, **kwargs)
 
-    def aggregate(self, *args, **kwargs):
+    def aggregate(self, *args, **kwargs) -> None:
         return MissingPandasLikeSeriesGroupBy.aggregate(self, *args, **kwargs)
 
     def transform(self, func, *args, **kwargs) -> Series:

databricks/koalas/indexes.py

@@ -2391,10 +2391,10 @@ def __abs__(self):
     def _with_new_scol(self, scol: spark.Column):
         raise NotImplementedError("Not supported for type MultiIndex")
 
-    def any(self, *args, **kwargs):
+    def any(self, *args, **kwargs) -> None:
         raise TypeError("cannot perform any with this index type: MultiIndex")
 
-    def all(self, *args, **kwargs):
+    def all(self, *args, **kwargs) -> None:
         raise TypeError("cannot perform all with this index type: MultiIndex")
 
     @staticmethod
@@ -2861,7 +2861,7 @@ def toPandas(self) -> pd.MultiIndex:
 
     toPandas.__doc__ = to_pandas.__doc__
 
-    def nunique(self, dropna=True):
+    def nunique(self, dropna=True) -> None:  # type: ignore
         raise NotImplementedError("nunique is not defined for MultiIndex")
 
     # TODO: add 'name' parameter after pd.MultiIndex.name is implemented
@@ -3074,13 +3074,13 @@ def value_counts(
 
     value_counts.__doc__ = IndexOpsMixin.value_counts.__doc__
 
-    def argmax(self):
+    def argmax(self) -> None:
         raise TypeError("reduction operation 'argmax' not allowed for this dtype")
 
-    def argmin(self):
+    def argmin(self) -> None:
         raise TypeError("reduction operation 'argmin' not allowed for this dtype")
 
-    def asof(self, label):
+    def asof(self, label) -> None:
         raise NotImplementedError(
             "only the default get_loc method is currently supported for MultiIndex"
         )