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.