TYP: misc return type annotations (#47558)

Author: twoertwein

pandas/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa
+from __future__ import annotations
 
 __docformat__ = "restructuredtext"
 
@@ -185,7 +186,7 @@
 __deprecated_num_index_names = ["Float64Index", "Int64Index", "UInt64Index"]
 
 
-def __dir__():
+def __dir__() -> list[str]:
     # GH43028
     # Int64Index etc. are deprecated, but we still want them to be available in the dir.
     # Remove in Pandas 2.0, when we remove Int64Index etc. from the code base.

pandas/_config/config.py

@@ -60,6 +60,7 @@
     Callable,
     Generic,
     Iterable,
+    Iterator,
     NamedTuple,
     cast,
 )
@@ -435,13 +436,13 @@ def __init__(self, *args) -> None:
 
         self.ops = list(zip(args[::2], args[1::2]))
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         self.undo = [(pat, _get_option(pat, silent=True)) for pat, val in self.ops]
 
         for pat, val in self.ops:
             _set_option(pat, val, silent=True)
 
-    def __exit__(self, *args):
+    def __exit__(self, *args) -> None:
         if self.undo:
             for pat, val in self.undo:
                 _set_option(pat, val, silent=True)
@@ -733,7 +734,7 @@ def pp(name: str, ks: Iterable[str]) -> list[str]:
 
 
 @contextmanager
-def config_prefix(prefix):
+def config_prefix(prefix) -> Iterator[None]:
     """
     contextmanager for multiple invocations of API with a common prefix
 

pandas/_testing/__init__.py

@@ -238,15 +238,15 @@
 _testing_mode_warnings = (DeprecationWarning, ResourceWarning)
 
 
-def set_testing_mode():
+def set_testing_mode() -> None:
     # set the testing mode filters
     testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None")
     if "deprecate" in testing_mode:
         for category in _testing_mode_warnings:
             warnings.simplefilter("always", category)
 
 
-def reset_testing_mode():
+def reset_testing_mode() -> None:
     # reset the testing mode filters
     testing_mode = os.environ.get("PANDAS_TESTING_MODE", "None")
     if "deprecate" in testing_mode:
@@ -257,7 +257,7 @@ def reset_testing_mode():
 set_testing_mode()
 
 
-def reset_display_options():
+def reset_display_options() -> None:
     """
     Reset the display options for printing and representing objects.
     """
@@ -333,38 +333,38 @@ def to_array(obj):
 # Others
 
 
-def getCols(k):
+def getCols(k) -> str:
     return string.ascii_uppercase[:k]
 
 
 # make index
-def makeStringIndex(k=10, name=None):
+def makeStringIndex(k=10, name=None) -> Index:
     return Index(rands_array(nchars=10, size=k), name=name)
 
 
-def makeCategoricalIndex(k=10, n=3, name=None, **kwargs):
+def makeCategoricalIndex(k=10, n=3, name=None, **kwargs) -> CategoricalIndex:
     """make a length k index or n categories"""
     x = rands_array(nchars=4, size=n, replace=False)
     return CategoricalIndex(
         Categorical.from_codes(np.arange(k) % n, categories=x), name=name, **kwargs
     )
 
 
-def makeIntervalIndex(k=10, name=None, **kwargs):
+def makeIntervalIndex(k=10, name=None, **kwargs) -> IntervalIndex:
     """make a length k IntervalIndex"""
     x = np.linspace(0, 100, num=(k + 1))
     return IntervalIndex.from_breaks(x, name=name, **kwargs)
 
 
-def makeBoolIndex(k=10, name=None):
+def makeBoolIndex(k=10, name=None) -> Index:
     if k == 1:
         return Index([True], name=name)
     elif k == 2:
         return Index([False, True], name=name)
     return Index([False, True] + [False] * (k - 2), name=name)
 
 
-def makeNumericIndex(k=10, name=None, *, dtype):
+def makeNumericIndex(k=10, name=None, *, dtype) -> NumericIndex:
     dtype = pandas_dtype(dtype)
     assert isinstance(dtype, np.dtype)
 
@@ -382,21 +382,21 @@ def makeNumericIndex(k=10, name=None, *, dtype):
     return NumericIndex(values, dtype=dtype, name=name)
 
 
-def makeIntIndex(k=10, name=None):
+def makeIntIndex(k=10, name=None) -> Int64Index:
     base_idx = makeNumericIndex(k, name=name, dtype="int64")
     return Int64Index(base_idx)
 
 
-def makeUIntIndex(k=10, name=None):
+def makeUIntIndex(k=10, name=None) -> UInt64Index:
     base_idx = makeNumericIndex(k, name=name, dtype="uint64")
     return UInt64Index(base_idx)
 
 
-def makeRangeIndex(k=10, name=None, **kwargs):
+def makeRangeIndex(k=10, name=None, **kwargs) -> RangeIndex:
     return RangeIndex(0, k, 1, name=name, **kwargs)
 
 
-def makeFloatIndex(k=10, name=None):
+def makeFloatIndex(k=10, name=None) -> Float64Index:
     base_idx = makeNumericIndex(k, name=name, dtype="float64")
     return Float64Index(base_idx)
 
@@ -456,57 +456,57 @@ def all_timeseries_index_generator(k: int = 10) -> Iterable[Index]:
 
 
 # make series
-def make_rand_series(name=None, dtype=np.float64):
+def make_rand_series(name=None, dtype=np.float64) -> Series:
     index = makeStringIndex(_N)
     data = np.random.randn(_N)
     data = data.astype(dtype, copy=False)
     return Series(data, index=index, name=name)
 
 
-def makeFloatSeries(name=None):
+def makeFloatSeries(name=None) -> Series:
     return make_rand_series(name=name)
 
 
-def makeStringSeries(name=None):
+def makeStringSeries(name=None) -> Series:
     return make_rand_series(name=name)
 
 
-def makeObjectSeries(name=None):
+def makeObjectSeries(name=None) -> Series:
     data = makeStringIndex(_N)
     data = Index(data, dtype=object)
     index = makeStringIndex(_N)
     return Series(data, index=index, name=name)
 
 
-def getSeriesData():
+def getSeriesData() -> dict[str, Series]:
     index = makeStringIndex(_N)
     return {c: Series(np.random.randn(_N), index=index) for c in getCols(_K)}
 
 
-def makeTimeSeries(nper=None, freq="B", name=None):
+def makeTimeSeries(nper=None, freq="B", name=None) -> Series:
     if nper is None:
         nper = _N
     return Series(
         np.random.randn(nper), index=makeDateIndex(nper, freq=freq), name=name
     )
 
 
-def makePeriodSeries(nper=None, name=None):
+def makePeriodSeries(nper=None, name=None) -> Series:
     if nper is None:
         nper = _N
     return Series(np.random.randn(nper), index=makePeriodIndex(nper), name=name)
 
 
-def getTimeSeriesData(nper=None, freq="B"):
+def getTimeSeriesData(nper=None, freq="B") -> dict[str, Series]:
     return {c: makeTimeSeries(nper, freq) for c in getCols(_K)}
 
 
-def getPeriodData(nper=None):
+def getPeriodData(nper=None) -> dict[str, Series]:
     return {c: makePeriodSeries(nper) for c in getCols(_K)}
 
 
 # make frame
-def makeTimeDataFrame(nper=None, freq="B"):
+def makeTimeDataFrame(nper=None, freq="B") -> DataFrame:
     data = getTimeSeriesData(nper, freq)
     return DataFrame(data)
 
@@ -533,14 +533,19 @@ def makeMixedDataFrame():
     return DataFrame(getMixedTypeDict()[1])
 
 
-def makePeriodFrame(nper=None):
+def makePeriodFrame(nper=None) -> DataFrame:
     data = getPeriodData(nper)
     return DataFrame(data)
 
 
 def makeCustomIndex(
-    nentries, nlevels, prefix="#", names=False, ndupe_l=None, idx_type=None
-):
+    nentries,
+    nlevels,
+    prefix="#",
+    names: bool | str | list[str] | None = False,
+    ndupe_l=None,
+    idx_type=None,
+) -> Index:
     """
     Create an index/multindex with given dimensions, levels, names, etc'
 
@@ -637,7 +642,8 @@ def keyfunc(x):
     # convert tuples to index
     if nentries == 1:
         # we have a single level of tuples, i.e. a regular Index
-        index = Index(tuples[0], name=names[0])
+        name = None if names is None else names[0]
+        index = Index(tuples[0], name=name)
     elif nlevels == 1:
         name = None if names is None else names[0]
         index = Index((x[0] for x in tuples), name=name)
@@ -659,7 +665,7 @@ def makeCustomDataframe(
     dtype=None,
     c_idx_type=None,
     r_idx_type=None,
-):
+) -> DataFrame:
     """
     Create a DataFrame using supplied parameters.
 
@@ -780,7 +786,7 @@ def _gen_unique_rand(rng, _extra_size):
     return i.tolist(), j.tolist()
 
 
-def makeMissingDataframe(density=0.9, random_state=None):
+def makeMissingDataframe(density=0.9, random_state=None) -> DataFrame:
     df = makeDataFrame()
     i, j = _create_missing_idx(*df.shape, density=density, random_state=random_state)
     df.values[i, j] = np.nan
@@ -854,7 +860,7 @@ def skipna_wrapper(x):
     return skipna_wrapper
 
 
-def convert_rows_list_to_csv_str(rows_list: list[str]):
+def convert_rows_list_to_csv_str(rows_list: list[str]) -> str:
     """
     Convert list of CSV rows to single CSV-formatted string for current OS.
 

pandas/_testing/_io.py

@@ -250,7 +250,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def can_connect(url, error_classes=None):
+def can_connect(url, error_classes=None) -> bool:
     """
     Try to connect to the given url. True if succeeds, False if OSError
     raised
@@ -424,7 +424,7 @@ def write_to_compressed(compression, path, data, dest="test"):
 # Plotting
 
 
-def close(fignum=None):
+def close(fignum=None) -> None:
     from matplotlib.pyplot import (
         close as _close,
         get_fignums,

pandas/_testing/_random.py

@@ -14,7 +14,7 @@ def randbool(size=(), p: float = 0.5):
 )
 
 
-def rands_array(nchars, size, dtype="O", replace=True):
+def rands_array(nchars, size, dtype="O", replace=True) -> np.ndarray:
     """
     Generate an array of byte strings.
     """
@@ -26,7 +26,7 @@ def rands_array(nchars, size, dtype="O", replace=True):
     return retval.astype(dtype)
 
 
-def rands(nchars):
+def rands(nchars) -> str:
     """
     Generate one random byte string.