From d44f6c11909616209ae14dd5805c378f6abbef19 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Mon, 4 Dec 2023 11:11:53 -0800 Subject: [PATCH] TST: dt64 units (#56261) * TST: dt64 units * fix on older pythons * typo fixup * mypy fixup * de-xfail --- pandas/core/dtypes/missing.py | 3 +- .../tests/frame/methods/test_reset_index.py | 10 +-- pandas/tests/indexes/interval/test_formats.py | 15 ++-- pandas/tests/io/excel/test_readers.py | 78 +++++++++++++------ pandas/tests/io/excel/test_writers.py | 39 ++++++++-- 5 files changed, 99 insertions(+), 46 deletions(-) diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py index a635ac77566e1..4dc0d477f89e8 100644 --- a/pandas/core/dtypes/missing.py +++ b/pandas/core/dtypes/missing.py @@ -686,7 +686,8 @@ def na_value_for_dtype(dtype: DtypeObj, compat: bool = True): if isinstance(dtype, ExtensionDtype): return dtype.na_value elif dtype.kind in "mM": - return dtype.type("NaT", "ns") + unit = np.datetime_data(dtype)[0] + return dtype.type("NaT", unit) elif dtype.kind == "f": return np.nan elif dtype.kind in "iu": diff --git a/pandas/tests/frame/methods/test_reset_index.py b/pandas/tests/frame/methods/test_reset_index.py index c989b3d26677c..20f0dcc816408 100644 --- a/pandas/tests/frame/methods/test_reset_index.py +++ b/pandas/tests/frame/methods/test_reset_index.py @@ -664,11 +664,8 @@ def test_reset_index_dtypes_on_empty_frame_with_multiindex(array, dtype): def test_reset_index_empty_frame_with_datetime64_multiindex(): # https://github.com/pandas-dev/pandas/issues/35606 - idx = MultiIndex( - levels=[[Timestamp("2020-07-20 00:00:00")], [3, 4]], - codes=[[], []], - names=["a", "b"], - ) + dti = pd.DatetimeIndex(["2020-07-20 00:00:00"], dtype="M8[ns]") + idx = MultiIndex.from_product([dti, [3, 4]], names=["a", "b"])[:0] df = DataFrame(index=idx, columns=["c", "d"]) result = df.reset_index() expected = DataFrame( @@ -681,7 +678,8 @@ def test_reset_index_empty_frame_with_datetime64_multiindex(): def test_reset_index_empty_frame_with_datetime64_multiindex_from_groupby(): # https://github.com/pandas-dev/pandas/issues/35657 - df = DataFrame({"c1": [10.0], "c2": ["a"], "c3": pd.to_datetime("2020-01-01")}) + dti = pd.DatetimeIndex(["2020-01-01"], dtype="M8[ns]") + df = DataFrame({"c1": [10.0], "c2": ["a"], "c3": dti}) df = df.head(0).groupby(["c2", "c3"])[["c1"]].sum() result = df.reset_index() expected = DataFrame( diff --git a/pandas/tests/indexes/interval/test_formats.py b/pandas/tests/indexes/interval/test_formats.py index bf335d154e186..edd7993056512 100644 --- a/pandas/tests/indexes/interval/test_formats.py +++ b/pandas/tests/indexes/interval/test_formats.py @@ -3,6 +3,7 @@ from pandas import ( DataFrame, + DatetimeIndex, Index, Interval, IntervalIndex, @@ -100,18 +101,14 @@ def test_get_values_for_csv(self, tuples, closed, expected_data): expected = np.array(expected_data) tm.assert_numpy_array_equal(result, expected) - def test_timestamp_with_timezone(self): + def test_timestamp_with_timezone(self, unit): # GH 55035 - index = IntervalIndex( - [ - Interval( - Timestamp("2020-01-01", tz="UTC"), Timestamp("2020-01-02", tz="UTC") - ) - ] - ) + left = DatetimeIndex(["2020-01-01"], dtype=f"M8[{unit}, UTC]") + right = DatetimeIndex(["2020-01-02"], dtype=f"M8[{unit}, UTC]") + index = IntervalIndex.from_arrays(left, right) result = repr(index) expected = ( "IntervalIndex([(2020-01-01 00:00:00+00:00, 2020-01-02 00:00:00+00:00]], " - "dtype='interval[datetime64[ns, UTC], right]')" + f"dtype='interval[datetime64[{unit}, UTC], right]')" ) assert result == expected diff --git a/pandas/tests/io/excel/test_readers.py b/pandas/tests/io/excel/test_readers.py index a9e83f6ef87a3..fe9ad8de5765f 100644 --- a/pandas/tests/io/excel/test_readers.py +++ b/pandas/tests/io/excel/test_readers.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import ( datetime, time, @@ -130,8 +132,15 @@ def df_ref(datapath): return df_ref -def adjust_expected(expected: DataFrame, read_ext: str) -> None: +def get_exp_unit(read_ext: str, engine: str | None) -> str: + return "ns" + + +def adjust_expected(expected: DataFrame, read_ext: str, engine: str) -> None: expected.index.name = None + unit = get_exp_unit(read_ext, engine) + # error: "Index" has no attribute "as_unit" + expected.index = expected.index.as_unit(unit) # type: ignore[attr-defined] def xfail_datetimes_with_pyxlsb(engine, request): @@ -225,7 +234,7 @@ def test_usecols_list(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref[["B", "C"]] - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) df1 = pd.read_excel( "test1" + read_ext, sheet_name="Sheet1", index_col=0, usecols=[0, 2, 3] @@ -246,7 +255,7 @@ def test_usecols_str(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref[["A", "B", "C"]] - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) df2 = pd.read_excel( "test1" + read_ext, sheet_name="Sheet1", index_col=0, usecols="A:D" @@ -264,7 +273,7 @@ def test_usecols_str(self, request, engine, read_ext, df_ref): tm.assert_frame_equal(df3, expected) expected = df_ref[["B", "C"]] - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) df2 = pd.read_excel( "test1" + read_ext, sheet_name="Sheet1", index_col=0, usecols="A,C,D" @@ -302,7 +311,7 @@ def test_usecols_diff_positional_int_columns_order( xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref[["A", "C"]] - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) result = pd.read_excel( "test1" + read_ext, sheet_name="Sheet1", index_col=0, usecols=usecols @@ -321,7 +330,7 @@ def test_read_excel_without_slicing(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) result = pd.read_excel("test1" + read_ext, sheet_name="Sheet1", index_col=0) tm.assert_frame_equal(result, expected) @@ -330,7 +339,7 @@ def test_usecols_excel_range_str(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref[["C", "D"]] - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) result = pd.read_excel( "test1" + read_ext, sheet_name="Sheet1", index_col=0, usecols="A,D:E" @@ -428,7 +437,7 @@ def test_excel_table(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) df1 = pd.read_excel("test1" + read_ext, sheet_name="Sheet1", index_col=0) df2 = pd.read_excel( @@ -446,6 +455,7 @@ def test_excel_table(self, request, engine, read_ext, df_ref): def test_reader_special_dtypes(self, request, engine, read_ext): xfail_datetimes_with_pyxlsb(engine, request) + unit = get_exp_unit(read_ext, engine) expected = DataFrame.from_dict( { "IntCol": [1, 2, -3, 4, 0], @@ -453,13 +463,16 @@ def test_reader_special_dtypes(self, request, engine, read_ext): "BoolCol": [True, False, True, True, False], "StrCol": [1, 2, 3, 4, 5], "Str2Col": ["a", 3, "c", "d", "e"], - "DateCol": [ - datetime(2013, 10, 30), - datetime(2013, 10, 31), - datetime(1905, 1, 1), - datetime(2013, 12, 14), - datetime(2015, 3, 14), - ], + "DateCol": Index( + [ + datetime(2013, 10, 30), + datetime(2013, 10, 31), + datetime(1905, 1, 1), + datetime(2013, 12, 14), + datetime(2015, 3, 14), + ], + dtype=f"M8[{unit}]", + ), }, ) basename = "test_types" @@ -578,7 +591,7 @@ def test_reader_dtype_str(self, read_ext, dtype, expected): actual = pd.read_excel(basename + read_ext, dtype=dtype) tm.assert_frame_equal(actual, expected) - def test_dtype_backend(self, read_ext, dtype_backend): + def test_dtype_backend(self, read_ext, dtype_backend, engine): # GH#36712 if read_ext in (".xlsb", ".xls"): pytest.skip(f"No engine for filetype: '{read_ext}'") @@ -621,6 +634,9 @@ def test_dtype_backend(self, read_ext, dtype_backend): expected["j"] = ArrowExtensionArray(pa.array([None, None])) else: expected = df + unit = get_exp_unit(read_ext, engine) + expected["i"] = expected["i"].astype(f"M8[{unit}]") + tm.assert_frame_equal(result, expected) def test_dtype_backend_and_dtype(self, read_ext): @@ -812,7 +828,7 @@ def test_sheet_name(self, request, read_ext, engine, df_ref): sheet_name = "Sheet1" expected = df_ref - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) df1 = pd.read_excel( filename + read_ext, sheet_name=sheet_name, index_col=0 @@ -1010,6 +1026,8 @@ def test_read_excel_multiindex(self, request, engine, read_ext): # see gh-4679 xfail_datetimes_with_pyxlsb(engine, request) + unit = get_exp_unit(read_ext, engine) + mi = MultiIndex.from_product([["foo", "bar"], ["a", "b"]]) mi_file = "testmultiindex" + read_ext @@ -1023,6 +1041,7 @@ def test_read_excel_multiindex(self, request, engine, read_ext): ], columns=mi, ) + expected[mi[2]] = expected[mi[2]].astype(f"M8[{unit}]") actual = pd.read_excel( mi_file, sheet_name="mi_column", header=[0, 1], index_col=0 @@ -1102,6 +1121,9 @@ def test_read_excel_multiindex_blank_after_name( mi_file = "testmultiindex" + read_ext mi = MultiIndex.from_product([["foo", "bar"], ["a", "b"]], names=["c1", "c2"]) + + unit = get_exp_unit(read_ext, engine) + expected = DataFrame( [ [1, 2.5, pd.Timestamp("2015-01-01"), True], @@ -1115,6 +1137,7 @@ def test_read_excel_multiindex_blank_after_name( names=["ilvl1", "ilvl2"], ), ) + expected[mi[2]] = expected[mi[2]].astype(f"M8[{unit}]") result = pd.read_excel( mi_file, sheet_name=sheet_name, @@ -1218,6 +1241,8 @@ def test_read_excel_skiprows(self, request, engine, read_ext): # GH 4903 xfail_datetimes_with_pyxlsb(engine, request) + unit = get_exp_unit(read_ext, engine) + actual = pd.read_excel( "testskiprows" + read_ext, sheet_name="skiprows_list", skiprows=[0, 2] ) @@ -1230,6 +1255,7 @@ def test_read_excel_skiprows(self, request, engine, read_ext): ], columns=["a", "b", "c", "d"], ) + expected["c"] = expected["c"].astype(f"M8[{unit}]") tm.assert_frame_equal(actual, expected) actual = pd.read_excel( @@ -1262,11 +1288,13 @@ def test_read_excel_skiprows(self, request, engine, read_ext): ], columns=["a", "b", "c", "d"], ) + expected["c"] = expected["c"].astype(f"M8[{unit}]") tm.assert_frame_equal(actual, expected) def test_read_excel_skiprows_callable_not_in(self, request, engine, read_ext): # GH 4903 xfail_datetimes_with_pyxlsb(engine, request) + unit = get_exp_unit(read_ext, engine) actual = pd.read_excel( "testskiprows" + read_ext, @@ -1282,6 +1310,7 @@ def test_read_excel_skiprows_callable_not_in(self, request, engine, read_ext): ], columns=["a", "b", "c", "d"], ) + expected["c"] = expected["c"].astype(f"M8[{unit}]") tm.assert_frame_equal(actual, expected) def test_read_excel_nrows(self, read_ext): @@ -1538,7 +1567,7 @@ def test_excel_table_sheet_by_index(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) with pd.ExcelFile("test1" + read_ext) as excel: df1 = pd.read_excel(excel, sheet_name=0, index_col=0) @@ -1565,7 +1594,7 @@ def test_sheet_name(self, request, engine, read_ext, df_ref): xfail_datetimes_with_pyxlsb(engine, request) expected = df_ref - adjust_expected(expected, read_ext) + adjust_expected(expected, read_ext, engine) filename = "test1" sheet_name = "Sheet1" @@ -1657,11 +1686,14 @@ def test_read_datetime_multiindex(self, request, engine, read_ext): f = "test_datetime_mi" + read_ext with pd.ExcelFile(f) as excel: actual = pd.read_excel(excel, header=[0, 1], index_col=0, engine=engine) - expected_column_index = MultiIndex.from_tuples( - [(pd.to_datetime("02/29/2020"), pd.to_datetime("03/01/2020"))], + + unit = get_exp_unit(read_ext, engine) + dti = pd.DatetimeIndex(["2020-02-29", "2020-03-01"], dtype=f"M8[{unit}]") + expected_column_index = MultiIndex.from_arrays( + [dti[:1], dti[1:]], names=[ - pd.to_datetime("02/29/2020").to_pydatetime(), - pd.to_datetime("03/01/2020").to_pydatetime(), + dti[0].to_pydatetime(), + dti[1].to_pydatetime(), ], ) expected = DataFrame([], index=[], columns=expected_column_index) diff --git a/pandas/tests/io/excel/test_writers.py b/pandas/tests/io/excel/test_writers.py index 74286a3ddd8ed..8452ec01a0936 100644 --- a/pandas/tests/io/excel/test_writers.py +++ b/pandas/tests/io/excel/test_writers.py @@ -35,6 +35,10 @@ from pandas.io.excel._util import _writers +def get_exp_unit(path: str) -> str: + return "ns" + + @pytest.fixture def frame(float_frame): """ @@ -461,6 +465,7 @@ def test_mixed(self, frame, path): tm.assert_frame_equal(mixed_frame, recons) def test_ts_frame(self, path): + unit = get_exp_unit(path) df = DataFrame( np.random.default_rng(2).standard_normal((5, 4)), columns=Index(list("ABCD"), dtype=object), @@ -471,10 +476,13 @@ def test_ts_frame(self, path): index = pd.DatetimeIndex(np.asarray(df.index), freq=None) df.index = index + expected = df[:] + expected.index = expected.index.as_unit(unit) + df.to_excel(path, sheet_name="test1") with ExcelFile(path) as reader: recons = pd.read_excel(reader, sheet_name="test1", index_col=0) - tm.assert_frame_equal(df, recons) + tm.assert_frame_equal(expected, recons) def test_basics_with_nan(self, frame, path): frame = frame.copy() @@ -538,6 +546,7 @@ def test_inf_roundtrip(self, path): def test_sheets(self, frame, path): # freq doesn't round-trip + unit = get_exp_unit(path) tsframe = DataFrame( np.random.default_rng(2).standard_normal((5, 4)), columns=Index(list("ABCD"), dtype=object), @@ -546,6 +555,9 @@ def test_sheets(self, frame, path): index = pd.DatetimeIndex(np.asarray(tsframe.index), freq=None) tsframe.index = index + expected = tsframe[:] + expected.index = expected.index.as_unit(unit) + frame = frame.copy() frame.iloc[:5, frame.columns.get_loc("A")] = np.nan @@ -562,7 +574,7 @@ def test_sheets(self, frame, path): recons = pd.read_excel(reader, sheet_name="test1", index_col=0) tm.assert_frame_equal(frame, recons) recons = pd.read_excel(reader, sheet_name="test2", index_col=0) - tm.assert_frame_equal(tsframe, recons) + tm.assert_frame_equal(expected, recons) assert 2 == len(reader.sheet_names) assert "test1" == reader.sheet_names[0] assert "test2" == reader.sheet_names[1] @@ -660,6 +672,7 @@ def test_excel_roundtrip_indexname(self, merge_cells, path): def test_excel_roundtrip_datetime(self, merge_cells, path): # datetime.date, not sure what to test here exactly + unit = get_exp_unit(path) # freq does not round-trip tsframe = DataFrame( @@ -678,12 +691,16 @@ def test_excel_roundtrip_datetime(self, merge_cells, path): with ExcelFile(path) as reader: recons = pd.read_excel(reader, sheet_name="test1", index_col=0) - tm.assert_frame_equal(tsframe, recons) + expected = tsframe[:] + expected.index = expected.index.as_unit(unit) + tm.assert_frame_equal(expected, recons) def test_excel_date_datetime_format(self, ext, path): # see gh-4133 # # Excel output format strings + unit = get_exp_unit(path) + df = DataFrame( [ [date(2014, 1, 31), date(1999, 9, 24)], @@ -700,6 +717,7 @@ def test_excel_date_datetime_format(self, ext, path): index=["DATE", "DATETIME"], columns=["X", "Y"], ) + df_expected = df_expected.astype(f"M8[{unit}]") with tm.ensure_clean(ext) as filename2: with ExcelWriter(path) as writer1: @@ -854,15 +872,20 @@ def test_to_excel_multiindex_cols(self, merge_cells, frame, path): def test_to_excel_multiindex_dates(self, merge_cells, path): # try multiindex with dates + unit = get_exp_unit(path) tsframe = DataFrame( np.random.default_rng(2).standard_normal((5, 4)), columns=Index(list("ABCD"), dtype=object), index=date_range("2000-01-01", periods=5, freq="B"), ) - new_index = [tsframe.index, np.arange(len(tsframe.index), dtype=np.int64)] - tsframe.index = MultiIndex.from_arrays(new_index) + tsframe.index = MultiIndex.from_arrays( + [ + tsframe.index.as_unit(unit), + np.arange(len(tsframe.index), dtype=np.int64), + ], + names=["time", "foo"], + ) - tsframe.index.names = ["time", "foo"] tsframe.to_excel(path, sheet_name="test1", merge_cells=merge_cells) with ExcelFile(path) as reader: recons = pd.read_excel(reader, sheet_name="test1", index_col=[0, 1]) @@ -1147,6 +1170,7 @@ def test_comment_empty_line(self, path): def test_datetimes(self, path): # Test writing and reading datetimes. For issue #9139. (xref #9185) + unit = get_exp_unit(path) datetimes = [ datetime(2013, 1, 13, 1, 2, 3), datetime(2013, 1, 13, 2, 45, 56), @@ -1165,7 +1189,8 @@ def test_datetimes(self, path): write_frame.to_excel(path, sheet_name="Sheet1") read_frame = pd.read_excel(path, sheet_name="Sheet1", header=0) - tm.assert_series_equal(write_frame["A"], read_frame["A"]) + expected = write_frame.astype(f"M8[{unit}]") + tm.assert_series_equal(expected["A"], read_frame["A"]) def test_bytes_io(self, engine): # see gh-7074