1import numpy as np
2
3from pandas import Categorical, DataFrame, Index, Series, Timestamp
4import pandas._testing as tm
5from pandas.core.arrays import IntervalArray, integer_array
6
7
8class TestGetNumericData:
9    def test_get_numeric_data_preserve_dtype(self):
10        # get the numeric data
11        obj = DataFrame({"A": [1, "2", 3.0]})
12        result = obj._get_numeric_data()
13        expected = DataFrame(index=[0, 1, 2], dtype=object)
14        tm.assert_frame_equal(result, expected)
15
16    def test_get_numeric_data(self):
17
18        datetime64name = np.dtype("M8[ns]").name
19        objectname = np.dtype(np.object_).name
20
21        df = DataFrame(
22            {"a": 1.0, "b": 2, "c": "foo", "f": Timestamp("20010102")},
23            index=np.arange(10),
24        )
25        result = df.dtypes
26        expected = Series(
27            [
28                np.dtype("float64"),
29                np.dtype("int64"),
30                np.dtype(objectname),
31                np.dtype(datetime64name),
32            ],
33            index=["a", "b", "c", "f"],
34        )
35        tm.assert_series_equal(result, expected)
36
37        df = DataFrame(
38            {
39                "a": 1.0,
40                "b": 2,
41                "c": "foo",
42                "d": np.array([1.0] * 10, dtype="float32"),
43                "e": np.array([1] * 10, dtype="int32"),
44                "f": np.array([1] * 10, dtype="int16"),
45                "g": Timestamp("20010102"),
46            },
47            index=np.arange(10),
48        )
49
50        result = df._get_numeric_data()
51        expected = df.loc[:, ["a", "b", "d", "e", "f"]]
52        tm.assert_frame_equal(result, expected)
53
54        only_obj = df.loc[:, ["c", "g"]]
55        result = only_obj._get_numeric_data()
56        expected = df.loc[:, []]
57        tm.assert_frame_equal(result, expected)
58
59        df = DataFrame.from_dict({"a": [1, 2], "b": ["foo", "bar"], "c": [np.pi, np.e]})
60        result = df._get_numeric_data()
61        expected = DataFrame.from_dict({"a": [1, 2], "c": [np.pi, np.e]})
62        tm.assert_frame_equal(result, expected)
63
64        df = result.copy()
65        result = df._get_numeric_data()
66        expected = df
67        tm.assert_frame_equal(result, expected)
68
69    def test_get_numeric_data_mixed_dtype(self):
70        # numeric and object columns
71
72        df = DataFrame(
73            {
74                "a": [1, 2, 3],
75                "b": [True, False, True],
76                "c": ["foo", "bar", "baz"],
77                "d": [None, None, None],
78                "e": [3.14, 0.577, 2.773],
79            }
80        )
81        result = df._get_numeric_data()
82        tm.assert_index_equal(result.columns, Index(["a", "b", "e"]))
83
84    def test_get_numeric_data_extension_dtype(self):
85        # GH#22290
86        df = DataFrame(
87            {
88                "A": integer_array([-10, np.nan, 0, 10, 20, 30], dtype="Int64"),
89                "B": Categorical(list("abcabc")),
90                "C": integer_array([0, 1, 2, 3, np.nan, 5], dtype="UInt8"),
91                "D": IntervalArray.from_breaks(range(7)),
92            }
93        )
94        result = df._get_numeric_data()
95        expected = df.loc[:, ["A", "C"]]
96        tm.assert_frame_equal(result, expected)
97