From ba6f6bf1ca404aa9dfce62d4223cca9cb2d33d76 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 8 Oct 2020 12:14:59 -0700 Subject: [PATCH 1/3] TYP: generic, series, frame --- pandas/core/frame.py | 4 ++-- pandas/core/generic.py | 17 +++++++++++++---- pandas/core/series.py | 11 ++++++++--- setup.cfg | 3 --- 4 files changed, 23 insertions(+), 12 deletions(-) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 29c63b030cdd5..37b8d355f5a2e 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -9306,8 +9306,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 diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 338b45b5503dc..0e36c8a378c19 100644 --- a/pandas/core/generic.py +++ b/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,6 +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() + assert r is not None # for mypy if not gc.get_referents(r) or r.shape == self.shape: self._is_copy = None return @@ -6684,6 +6689,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 +7271,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] ) diff --git a/pandas/core/series.py b/pandas/core/series.py index 0852f1b650ae6..f9016b4463d6e 100644 --- a/pandas/core/series.py +++ b/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(): diff --git a/setup.cfg b/setup.cfg index 8d3d79789a252..023c7b1ffcad4 100644 --- a/setup.cfg +++ b/setup.cfg @@ -208,9 +208,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 From 5e3b2f84c06379d3c138b7c8c8a98e42d8c49808 Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 8 Oct 2020 13:18:34 -0700 Subject: [PATCH 2/3] fix assertion --- pandas/core/generic.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 0e36c8a378c19..8cc6ca6630099 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -3834,8 +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() - assert r is not None # for mypy - 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 From b3c9feb576c227f313dfe17ae5b871c0d17d1c60 Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 10 Oct 2020 08:23:06 -0700 Subject: [PATCH 3/3] test --- pandas/tests/series/methods/test_count.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pandas/tests/series/methods/test_count.py b/pandas/tests/series/methods/test_count.py index 1ca48eeb7c441..19290b6a5c23f 100644 --- a/pandas/tests/series/methods/test_count.py +++ b/pandas/tests/series/methods/test_count.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import pandas as pd from pandas import Categorical, MultiIndex, Series @@ -6,6 +7,13 @@ 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)