TYP: generic, series, frame (#36989)

Author: jbrockmendel

pandas/core/frame.py

@@ -9311,8 +9311,8 @@ def _AXIS_NAMES(self) -> Dict[int, str]:
 ops.add_special_arithmetic_methods(DataFrame)
 
 
-def _from_nested_dict(data):
-    new_data = collections.defaultdict(dict)
+def _from_nested_dict(data) -> collections.defaultdict:
+    new_data: collections.defaultdict = collections.defaultdict(dict)
     for index, s in data.items():
         for col, v in s.items():
             new_data[col][index] = v

pandas/core/generic.py

@@ -112,6 +112,7 @@
 if TYPE_CHECKING:
     from pandas._libs.tslibs import BaseOffset
 
+    from pandas.core.frame import DataFrame
     from pandas.core.resample import Resampler
     from pandas.core.series import Series
     from pandas.core.window.indexers import BaseIndexer
@@ -130,7 +131,7 @@
 )
 
 
-def _single_replace(self, to_replace, method, inplace, limit):
+def _single_replace(self: "Series", to_replace, method, inplace, limit):
     """
     Replaces values in a Series using the fill method specified when no
     replacement value is given in the replace method
@@ -541,6 +542,7 @@ def _get_cleaned_column_resolvers(self) -> Dict[str, ABCSeries]:
         from pandas.core.computation.parsing import clean_column_name
 
         if isinstance(self, ABCSeries):
+            self = cast("Series", self)
             return {clean_column_name(self.name): self}
 
         return {
@@ -1995,9 +1997,10 @@ def _repr_data_resource_(self):
         """
         if config.get_option("display.html.table_schema"):
             data = self.head(config.get_option("display.max_rows"))
-            payload = json.loads(
-                data.to_json(orient="table"), object_pairs_hook=collections.OrderedDict
-            )
+
+            as_json = data.to_json(orient="table")
+            as_json = cast(str, as_json)
+            payload = json.loads(as_json, object_pairs_hook=collections.OrderedDict)
             return payload
 
     # ----------------------------------------------------------------------
@@ -3113,6 +3116,7 @@ def to_latex(
         if multirow is None:
             multirow = config.get_option("display.latex.multirow")
 
+        self = cast("DataFrame", self)
         formatter = DataFrameFormatter(
             self,
             columns=columns,
@@ -3830,7 +3834,7 @@ def _check_setitem_copy(self, stacklevel=4, t="setting", force=False):
         # the copy weakref
         if self._is_copy is not None and not isinstance(self._is_copy, str):
             r = self._is_copy()
-            if not gc.get_referents(r) or r.shape == self.shape:
+            if not gc.get_referents(r) or (r is not None and r.shape == self.shape):
                 self._is_copy = None
                 return
 
@@ -6684,6 +6688,7 @@ def replace(
                     return self.apply(
                         _single_replace, args=(to_replace, method, inplace, limit)
                     )
+                self = cast("Series", self)
                 return _single_replace(self, to_replace, method, inplace, limit)
 
             if not is_dict_like(to_replace):
@@ -7265,10 +7270,13 @@ def asof(self, where, subset=None):
         nulls = self.isna() if is_series else self[subset].isna().any(1)
         if nulls.all():
             if is_series:
+                self = cast("Series", self)
                 return self._constructor(np.nan, index=where, name=self.name)
             elif is_list:
+                self = cast("DataFrame", self)
                 return self._constructor(np.nan, index=where, columns=self.columns)
             else:
+                self = cast("DataFrame", self)
                 return self._constructor_sliced(
                     np.nan, index=self.columns, name=where[0]
                 )

pandas/core/series.py

@@ -1789,12 +1789,17 @@ def count(self, level=None):
         """
         if level is None:
             return notna(self.array).sum()
+        elif not isinstance(self.index, MultiIndex):
+            raise ValueError("Series.count level is only valid with a MultiIndex")
+
+        index = self.index
+        assert isinstance(index, MultiIndex)  # for mypy
 
         if isinstance(level, str):
-            level = self.index._get_level_number(level)
+            level = index._get_level_number(level)
 
-        lev = self.index.levels[level]
-        level_codes = np.array(self.index.codes[level], subok=False, copy=True)
+        lev = index.levels[level]
+        level_codes = np.array(index.codes[level], subok=False, copy=True)
 
         mask = level_codes == -1
         if mask.any():

pandas/tests/series/methods/test_count.py

@@ -1,11 +1,19 @@
 import numpy as np
+import pytest
 
 import pandas as pd
 from pandas import Categorical, MultiIndex, Series
 import pandas._testing as tm
 
 
 class TestSeriesCount:
+    def test_count_level_without_multiindex(self):
+        ser = pd.Series(range(3))
+
+        msg = "Series.count level is only valid with a MultiIndex"
+        with pytest.raises(ValueError, match=msg):
+            ser.count(level=1)
+
     def test_count(self, datetime_series):
         assert datetime_series.count() == len(datetime_series)
 

setup.cfg

@@ -205,9 +205,6 @@ check_untyped_defs=False
 [mypy-pandas.core.reshape.merge]
 check_untyped_defs=False
 
-[mypy-pandas.core.series]
-check_untyped_defs=False
-
 [mypy-pandas.core.window.common]
 check_untyped_defs=False