Add type hinting to public-facing API in pandas/core/generic.py #26792

pandas/core/generic.py

@@ -19,6 +19,7 @@
     Optional,
     Sequence,
     Set,
+    Tuple,
     Union,
 )
 import warnings
@@ -91,6 +92,7 @@
 from pandas.io.formats.printing import pprint_thing
 from pandas.tseries.frequencies import to_offset
 
+
 # goal is to be able to define the docs close to function, while still being
 # able to share
 _shared_docs = dict()  # type: Dict[str, str]
@@ -553,7 +555,7 @@ def _stat_axis(self):
         return getattr(self, self._stat_axis_name)
 
     @property
-    def shape(self):
+    def shape(self) -> Tuple[int]:
         """
         Return a tuple of axis dimensions
         """
@@ -569,7 +571,7 @@ def axes(self):
         return [self._get_axis(a) for a in self._AXIS_ORDERS]
 
     @property
-    def ndim(self):
+    def ndim(self) -> int:
         """
         Return an int representing the number of axes / array dimensions.
 
@@ -592,7 +594,7 @@ def ndim(self):
         return self._data.ndim
 
     @property
-    def size(self):
+    def size(self) -> int:
         """
         Return an int representing the number of elements in this object.
 
@@ -625,7 +627,7 @@ def _obj_with_exclusions(self):
         """ internal compat with SelectionMixin """
         return self
 
-    def set_axis(self, labels, axis=0, inplace=False):
+    def set_axis(self, labels: Sequence[Any], axis=0, inplace=False) -> "NDFrame":
         """
         Assign desired index to given axis.
 
@@ -728,7 +730,7 @@ def _set_axis(self, axis, labels):
         self._data.set_axis(axis, labels)
         self._clear_item_cache()
 
-    def transpose(self, *args, **kwargs):
+    def transpose(self, *args, **kwargs) -> "NDFrame":
         """
         Permute the dimensions of the %(klass)s
 
@@ -772,7 +774,7 @@ def transpose(self, *args, **kwargs):
         nv.validate_transpose(tuple(), kwargs)
         return self._constructor(new_values, **new_axes).__finalize__(self)
 
-    def swapaxes(self, axis1, axis2, copy=True):
+    def swapaxes(self, axis1: int, axis2: int, copy=True) -> "NDFrame":
         """
         Interchange axes and swap values axes appropriately.
 
@@ -797,7 +799,9 @@ def swapaxes(self, axis1, axis2, copy=True):
 
         return self._constructor(new_values, *new_axes).__finalize__(self)
 
-    def droplevel(self, level, axis=0):
+    def droplevel(
+        self, level: Union[int, str, Sequence[Union[int, str]]], axis=0
+    ) -> "NDFrame":
         """
         Return DataFrame with requested index / column level(s) removed.
 
@@ -856,7 +860,7 @@ def droplevel(self, level, axis=0):
         result = self.set_axis(new_labels, axis=axis, inplace=False)
         return result
 
-    def pop(self, item):
+    def pop(self, item: str) -> "pd.Series":
         """
         Return item and drop from frame. Raise KeyError if not found.
 
@@ -906,7 +910,7 @@ def pop(self, item):
 
         return result
 
-    def squeeze(self, axis=None):
+    def squeeze(self, axis=None) -> "NDFrame":
         """
         Squeeze 1 dimensional axis objects into scalars.
 
@@ -1017,7 +1021,7 @@ def squeeze(self, axis=None):
             )
         ]
 
-    def swaplevel(self, i=-2, j=-1, axis=0):
+    def swaplevel(self, i=-2, j=-1, axis=0) -> "NDFrame":
         """
         Swap levels i and j in a MultiIndex on a particular axis
 
@@ -1039,7 +1043,7 @@ def swaplevel(self, i=-2, j=-1, axis=0):
     # ----------------------------------------------------------------------
     # Rename
 
-    def rename(self, *args, **kwargs):
+    def rename(self, *args, **kwargs) -> "NDFrame":
         """
         Alter axes input function or functions. Function / dict values must be
         unique (1-to-1). Labels not contained in a dict / Series will be left
@@ -1204,7 +1208,7 @@ def rename(self, *args, **kwargs):
             return result.__finalize__(self)
 
     @rewrite_axis_style_signature("mapper", [("copy", True), ("inplace", False)])
-    def rename_axis(self, mapper=sentinel, **kwargs):
+    def rename_axis(self, mapper=sentinel, **kwargs) -> "NDFrame":
         """
         Set the name of the axis for the index or columns.
 
@@ -1374,7 +1378,9 @@ class  name
             if not inplace:
                 return result
 
-    def _set_axis_name(self, name, axis=0, inplace=False):
+    def _set_axis_name(
+        self, name: Union[str, List[str]], axis=0, inplace=False
+    ) -> "NDFrame":
         """
         Set the name(s) of the axis.
 
@@ -1443,7 +1449,7 @@ def _indexed_same(self, other):
             self._get_axis(a).equals(other._get_axis(a)) for a in self._AXIS_ORDERS
         )
 
-    def equals(self, other):
+    def equals(self, other: "NDFrame") -> bool:
         """
         Test whether two objects contain the same elements.
 
@@ -1583,7 +1589,7 @@ def __nonzero__(self):
 
     __bool__ = __nonzero__
 
-    def bool(self):
+    def bool(self) -> bool:
         """
         Return the bool of a single element PandasObject.
 
@@ -1621,7 +1627,7 @@ def __round__(self, decimals=0):
     # operations should utilize/extend these methods when possible so that we
     # have consistent precedence and validation logic throughout the library.
 
-    def _is_level_reference(self, key, axis=0):
+    def _is_level_reference(self, key: str, axis=0) -> bool:
         """
         Test whether a key is a level reference for a given axis.
 
@@ -1651,7 +1657,7 @@ def _is_level_reference(self, key, axis=0):
             and not self._is_label_reference(key, axis=axis)
         )
 
-    def _is_label_reference(self, key, axis=0):
+    def _is_label_reference(self, key: str, axis=0) -> bool:
         """
         Test whether a key is a label reference for a given axis.
 
@@ -1680,7 +1686,7 @@ def _is_label_reference(self, key, axis=0):
             and any(key in self.axes[ax] for ax in other_axes)
         )
 
-    def _is_label_or_level_reference(self, key, axis=0):
+    def _is_label_or_level_reference(self, key: str, axis=0) -> bool:
         """
         Test whether a key is a label or level reference for a given axis.
 
@@ -1704,7 +1710,7 @@ def _is_label_or_level_reference(self, key, axis=0):
             key, axis=axis
         )
 
-    def _check_label_or_level_ambiguity(self, key, axis=0):
+    def _check_label_or_level_ambiguity(self, key: str, axis=0) -> None:
         """
         Check whether `key` is ambiguous.
 
@@ -1753,7 +1759,7 @@ def _check_label_or_level_ambiguity(self, key, axis=0):
             )
             raise ValueError(msg)
 
-    def _get_label_or_level_values(self, key, axis=0):
+    def _get_label_or_level_values(self, key: str, axis=0) -> np.ndarray:
         """
         Return a 1-D array of values associated with `key`, a label or level
         from the given `axis`.
@@ -1825,7 +1831,7 @@ def _get_label_or_level_values(self, key, axis=0):
 
         return values
 
-    def _drop_labels_or_levels(self, keys, axis=0):
+    def _drop_labels_or_levels(self, keys: Union[str, List[str]], axis=0) -> "NDFrame":
         """
         Drop labels and/or levels for the given `axis`.
 
@@ -1961,7 +1967,7 @@ def __contains__(self, key):
         return key in self._info_axis
 
     @property
-    def empty(self):
+    def empty(self) -> bool:
         """
         Indicator whether DataFrame is empty.
 
@@ -2466,7 +2472,9 @@ def to_json(
             indent=indent,
         )
 
-    def to_hdf(self, path_or_buf, key, **kwargs):
+    def to_hdf(
+        self, path_or_buf: Union[str, FilePathOrBuffer], key: str, **kwargs
+    ) -> None:
         """
         Write the contained data to an HDF5 file using HDFStore.
 
@@ -2571,7 +2579,9 @@ def to_hdf(self, path_or_buf, key, **kwargs):
 
         pytables.to_hdf(path_or_buf, key, self, **kwargs)
 
-    def to_msgpack(self, path_or_buf=None, encoding="utf-8", **kwargs):
+    def to_msgpack(
+        self, path_or_buf: Optional[FilePathOrBuffer] = None, encoding="utf-8", **kwargs
+    ) -> None:
         """
         Serialize object to input file path using msgpack format.
 
@@ -2767,7 +2777,9 @@ def to_sql(
             method=method,
         )
 
-    def to_pickle(self, path, compression="infer", protocol=pickle.HIGHEST_PROTOCOL):
+    def to_pickle(
+        self, path: str, compression="infer", protocol=pickle.HIGHEST_PROTOCOL
+    ) -> None:
         """
         Pickle (serialize) object to file.
 
@@ -2823,7 +2835,7 @@ def to_pickle(self, path, compression="infer", protocol=pickle.HIGHEST_PROTOCOL)
 
         to_pickle(self, path, compression=compression, protocol=protocol)
 
-    def to_clipboard(self, excel=True, sep=None, **kwargs):
+    def to_clipboard(self, excel=True, sep=None, **kwargs) -> None:
         r"""
         Copy object to the system clipboard.