Annotate DataFrame (Part 1) (pandas-dev#26867)

Author: vaibhavhrt

pandas/_typing.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, AnyStr, Optional, TypeVar, Union
+from typing import IO, TYPE_CHECKING, AnyStr, Iterable, Optional, TypeVar, Union
 
 import numpy as np
 
@@ -29,5 +29,8 @@
 Axis = Union[str, int]
 Ordered = Optional[bool]
 
+# use Collection after we drop support for py35
+Axes = Iterable
+
 # to maintain type information across generic functions and parametrization
 _T = TypeVar("_T")

pandas/core/frame.py

@@ -80,6 +80,7 @@
 )
 from pandas.core.dtypes.missing import isna, notna
 
+from pandas._typing import Axes, Dtype
 from pandas.core import algorithms, common as com, nanops, ops
 from pandas.core.accessor import CachedAccessor
 from pandas.core.arrays import Categorical, ExtensionArray
@@ -370,7 +371,7 @@ class DataFrame(NDFrame):
     """
 
     @property
-    def _constructor(self):
+    def _constructor(self) -> Type["DataFrame"]:
         return DataFrame
 
     _constructor_sliced = Series  # type: Type[Series]
@@ -386,7 +387,14 @@ def _constructor_expanddim(self):
     # ----------------------------------------------------------------------
     # Constructors
 
-    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False):
+    def __init__(
+        self,
+        data=None,
+        index: Optional[Axes] = None,
+        columns: Optional[Axes] = None,
+        dtype: Optional[Dtype] = None,
+        copy: bool = False,
+    ):
         if data is None:
             data = {}
         if dtype is not None:
@@ -481,7 +489,7 @@ def __init__(self, data=None, index=None, columns=None, dtype=None, copy=False):
     # ----------------------------------------------------------------------
 
     @property
-    def axes(self):
+    def axes(self) -> List[Index]:
         """
         Return a list representing the axes of the DataFrame.
 
@@ -498,7 +506,7 @@ def axes(self):
         return [self.index, self.columns]
 
     @property
-    def shape(self):
+    def shape(self) -> Tuple[int, int]:
         """
         Return a tuple representing the dimensionality of the DataFrame.
 
@@ -520,7 +528,7 @@ def shape(self):
         return len(self.index), len(self.columns)
 
     @property
-    def _is_homogeneous_type(self):
+    def _is_homogeneous_type(self) -> bool:
         """
         Whether all the columns in a DataFrame have the same type.