Adress review

pandas-dev · jreback · Nov 4, 2020 · Oct 3, 2020 · Oct 3, 2020 · Oct 4, 2020
commit c93590b1bd6b8869433a2bd19dbdb48840339983
diff --git a/doc/source/whatsnew/v1.2.0.rst b/doc/source/whatsnew/v1.2.0.rst
@@ -410,7 +410,7 @@ Groupby/resample/rolling
 - Bug in :meth:`Rolling.sum()` returned wrong values when dtypes where mixed between float and integer and axis was equal to one (:issue:`20649`, :issue:`35596`)
 - Bug in :meth:`Rolling.count` returned ``np.nan`` with :class:`pandas.api.indexers.FixedForwardWindowIndexer` as window, ``min_periods=0`` and only missing values in window (:issue:`35579`)
 - Bug where :class:`pandas.core.window.Rolling` produces incorrect window sizes when using a ``PeriodIndex`` (:issue:`34225`)
-- Bug in :meth:`Groupby` dropped ``nan`` groups from result with ``dropna=False`` when grouping over a single column (:issue:`35646`, :issue:`35542`)
+- Bug in :meth:`DataFrame.groupby` dropped ``nan`` groups from result with ``dropna=False`` when grouping over a single column (:issue:`35646`, :issue:`35542`)
 
 Reshaping
 ^^^^^^^^^

diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
@@ -1298,35 +1298,12 @@ def test_groupby_nat_exclude():
 
 
 def test_groupby_two_group_keys_all_nan():
-    # Grouping over two group keys all nan raised an error previously
+    # GH #36842: Grouping over two group keys shouldn't raise an error
     df = pd.DataFrame({"a": [np.nan, np.nan], "b": [np.nan, np.nan], "c": [1, 2]})
     result = df.groupby(["a", "b"]).indices
     assert result == {}
 
 
-def test_groupby_nan_included():
-    # GH 35646, GH 35542
-    data = {"group": ["g1", np.nan, "g1", "g2", np.nan], "B": [0, 1, 2, 3, 4]}
-    df = pd.DataFrame(data)
-    grouped = df.groupby("group", dropna=False)
-    result = grouped.indices
-    dtype = "int64"
-    expected = {
-        "g1": np.array([0, 2], dtype=dtype),
-        "g2": np.array([3], dtype=dtype),
-        np.nan: np.array([1, 4], dtype=dtype),
-    }
-    for result_values, expected_values in zip(result.values(), expected.values()):
-        tm.assert_numpy_array_equal(result_values, expected_values)
-    assert np.isnan(list(result.keys())[2])
-
-    result = grouped.mean()
-    expected = pd.DataFrame(
-        {"B": [1.0, 3.0, 2.5]}, index=pd.Index(["g1", "g2", np.nan], name="group")
-    )
-    tm.assert_frame_equal(result, expected)
-
-
 def test_groupby_2d_malformed():
     d = DataFrame(index=range(2))
     d["group"] = ["g1", "g2"]

diff --git a/pandas/tests/groupby/test_groupby_dropna.py b/pandas/tests/groupby/test_groupby_dropna.py
@@ -2,7 +2,7 @@
 import pytest
 
 import pandas as pd
-import pandas.testing as tm
+import pandas._testing as tm
 
 
 @pytest.mark.parametrize(
@@ -336,3 +336,27 @@ def test_groupby_apply_with_dropna_for_multi_index(dropna, data, selected_data,
 
     expected = pd.DataFrame(selected_data, index=mi)
     tm.assert_frame_equal(result, expected)
+
+
+def test_groupby_nan_included():
+    # GH 35646, GH 35542
+    data = {"group": ["g1", np.nan, "g1", "g2", np.nan], "B": [0, 1, 2, 3, 4]}
+    df = pd.DataFrame(data)
+    grouped = df.groupby("group", dropna=False)
+    result = grouped.indices
+    dtype = "int64"
+    expected = {
+        "g1": np.array([0, 2], dtype=dtype),
+        "g2": np.array([3], dtype=dtype),
+        np.nan: np.array([1, 4], dtype=dtype),
+    }
+    for result_values, expected_values in zip(result.values(), expected.values()):
+        tm.assert_numpy_array_equal(result_values, expected_values)
+    assert np.isnan(list(result.keys())[2])
+    assert list(result.keys())[0:2] == ["g1", "g2"]
+
+    result = grouped.mean()
+    expected = pd.DataFrame(
+        {"B": [1.0, 3.0, 2.5]}, index=pd.Index(["g1", "g2", np.nan], name="group")
+    )
+    tm.assert_frame_equal(result, expected)