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