Add type annotations for Indexing

databricks/koalas/frame.py

@@ -9988,9 +9988,9 @@ class  max_speed
         if not is_list_like(indices) or isinstance(indices, (dict, set)):
             raise ValueError("`indices` must be a list-like except dict or set")
         if axis == 0:
-            return self.iloc[indices, :]
+            return self.iloc[indices, :]  # type: ignore
         else:
-            return self.iloc[:, indices]
+            return self.iloc[:, indices]  # type: ignore
 
     def eval(self, expr, inplace=False) -> Optional["DataFrame"]:
         """

databricks/koalas/generic.py

@@ -15,7 +15,7 @@
 #
 
 """
-A base class to be monkey-patched to DataFrame/Column to behave similar to pandas DataFrame/Series.
+A base class of DataFrame/Column to behave similar to pandas DataFrame/Series.
 """
 from abc import ABCMeta, abstractmethod
 from collections import Counter
@@ -2514,25 +2514,25 @@ def ffill(self, axis=None, inplace=False, limit=None):
     pad = ffill
 
     @property
-    def at(self):
+    def at(self) -> AtIndexer:
         return AtIndexer(self)
 
     at.__doc__ = AtIndexer.__doc__
 
     @property
-    def iat(self):
+    def iat(self) -> iAtIndexer:
         return iAtIndexer(self)
 
     iat.__doc__ = iAtIndexer.__doc__
 
     @property
-    def iloc(self):
+    def iloc(self) -> iLocIndexer:
         return iLocIndexer(self)
 
     iloc.__doc__ = iLocIndexer.__doc__
 
     @property
-    def loc(self):
+    def loc(self) -> LocIndexer:
         return LocIndexer(self)
 
     loc.__doc__ = LocIndexer.__doc__

databricks/koalas/indexing.py

@@ -21,7 +21,7 @@
 from collections import OrderedDict
 from collections.abc import Iterable
 from functools import reduce
-from typing import Any, Optional, List, Tuple, TYPE_CHECKING
+from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union
 
 from pandas.api.types import is_list_like
 from pyspark import sql as spark
@@ -30,12 +30,14 @@
 from pyspark.sql.utils import AnalysisException
 import numpy as np
 
+from databricks import koalas as ks
 from databricks.koalas.internal import (
     InternalFrame,
     NATURAL_ORDER_COLUMN_NAME,
     SPARK_DEFAULT_SERIES_NAME,
 )
 from databricks.koalas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError
+from databricks.koalas.typedef import Scalar
 from databricks.koalas.utils import (
     is_name_like_tuple,
     is_name_like_value,
@@ -123,7 +125,7 @@ class AtIndexer(IndexerLike):
     array([ 4, 20])
     """
 
-    def __getitem__(self, key):
+    def __getitem__(self, key) -> Union["ks.Series", "ks.DataFrame", Scalar]:
         if self._is_df:
             if not isinstance(key, tuple) or len(key) != 2:
                 raise TypeError("Use DataFrame.at like .at[row_index, column_name]")
@@ -212,7 +214,7 @@ class iAtIndexer(IndexerLike):
     2
     """
 
-    def __getitem__(self, key):
+    def __getitem__(self, key) -> Union["ks.Series", "ks.DataFrame", Scalar]:
         if self._is_df:
             if not isinstance(key, tuple) or len(key) != 2:
                 raise TypeError(
@@ -387,7 +389,7 @@ def _select_cols_else(
         """ Select columns by other type key. """
         pass
 
-    def __getitem__(self, key):
+    def __getitem__(self, key) -> Union["ks.Series", "ks.DataFrame"]:
         from databricks.koalas.frame import DataFrame
         from databricks.koalas.series import Series, first_series