Fixed arrays tests to check dtypes

Author: 5hirish

pandas/core/groupby/groupby.py

@@ -759,20 +759,19 @@ def _try_cast(self, result, obj, numeric_only=False):
         else:
             dtype = obj.dtype
 
-        if is_extension_array_dtype(dtype):
-            # The function can return something of any type, so check
-            # if the type is compatible with the calling EA.
-            try:
-                result = obj.values._from_sequence(result)
-            except Exception:
-                # https://github.com/pandas-dev/pandas/issues/22850
-                # pandas has no control over what 3rd-party ExtensionArrays
-                # do in _values_from_sequence. We still want ops to work
-                # though, so we catch any regular Exception.
-                pass
-
         if not is_scalar(result):
-            if numeric_only and is_numeric_dtype(dtype) or not numeric_only:
+            if is_extension_array_dtype(dtype):
+                # The function can return something of any type, so check
+                # if the type is compatible with the calling EA.
+                try:
+                    result = obj.values._from_sequence(result)
+                except Exception:
+                    # https://github.com/pandas-dev/pandas/issues/22850
+                    # pandas has no control over what 3rd-party ExtensionArrays
+                    # do in _values_from_sequence. We still want ops to work
+                    # though, so we catch any regular Exception.
+                    pass
+            elif numeric_only and is_numeric_dtype(dtype) or not numeric_only:
                 result = maybe_downcast_to_dtype(result, dtype)
 
         return result

pandas/tests/arrays/test_integer.py

@@ -651,11 +651,9 @@ def test_preserve_dtypes(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="int64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
+
+    assert result.dtypes['B'].name == 'float64'
+    assert result.dtypes['C'].name == 'Int64'
 
 
 @pytest.mark.parametrize('op', ['mean'])
@@ -674,11 +672,23 @@ def test_reduce_to_float(op):
 
     # groupby
     result = getattr(df.groupby("A"), op)()
-    expected = pd.DataFrame({
-        "B": np.array([1.0, 3.0]),
-        "C": np.array([1, 3], dtype="float64")
-    }, index=pd.Index(['a', 'b'], name='A'))
-    tm.assert_frame_equal(result, expected)
+
+    assert result.dtypes['B'].name == 'float64'
+    assert result.dtypes['C'].name == 'Int64'
+
+
+@pytest.mark.parametrize('op', ['sum'])
+def test_groupby_extension_array(op):
+    # GH23227
+    # groupby on an extension array should return the extension array type
+    df = pd.DataFrame({
+        'Int': pd.Series([1, 2, 3], dtype='Int64'),
+        'A': [1, 2, 1]
+    })
+
+    result = getattr(df.groupby('A').Int, op)()
+    assert result is not None
+    assert result.dtype.name == 'Int64'
 
 
 def test_astype_nansafe():

pandas/tests/groupby/test_grouping.py

@@ -282,16 +282,6 @@ def test_groupby_categorical_index_and_columns(self, observed):
         expected = DataFrame(data=expected_data.T, index=expected_columns)
         assert_frame_equal(result, expected)
 
-    def test_groupby_extension_array(self):
-
-        # GH23227
-        # groupby on an extension array should return the extension array type
-        df = pd.DataFrame({'Int': pd.Series([1, 2, 3], dtype='Int64'),
-                           'A': [1, 2, 1]})
-        result = df.groupby('A').Int.sum()
-        assert result is not None
-        assert result.dtype.name == 'Int64'
-
     def test_grouper_getting_correct_binner(self):
 
         # GH 10063

pandas/tests/sparse/test_groupby.py

@@ -55,5 +55,5 @@ def test_groupby_includes_fill_value(fill_value):
     sdf = df.to_sparse(fill_value=fill_value)
     result = sdf.groupby('a').sum()
     expected = df.groupby('a').sum()
-    tm.assert_frame_equal(result, expected,
+    tm.assert_frame_equal(result, expected.to_sparse(fill_value=fill_value),
                           check_index_type=False)