TYP: dt64 unit in pytables tests (#56118)

Author: jbrockmendel

Diff for: pandas/tests/io/pytables/test_put.py

@@ -1,4 +1,3 @@
-import datetime
 import re
 
 import numpy as np
@@ -15,6 +14,7 @@
     Series,
     _testing as tm,
     concat,
+    date_range,
 )
 from pandas.tests.io.pytables.common import (
     _maybe_remove,
@@ -279,15 +279,9 @@ def test_store_multiindex(setup_path):
     with ensure_clean_store(setup_path) as store:
 
         def make_index(names=None):
-            return MultiIndex.from_tuples(
-                [
-                    (datetime.datetime(2013, 12, d), s, t)
-                    for d in range(1, 3)
-                    for s in range(2)
-                    for t in range(3)
-                ],
-                names=names,
-            )
+            dti = date_range("2013-12-01", "2013-12-02")
+            mi = MultiIndex.from_product([dti, range(2), range(3)], names=names)
+            return mi
 
         # no names
         _maybe_remove(store, "df")
@@ -306,11 +300,11 @@ def make_index(names=None):
         tm.assert_frame_equal(store.select("df"), df)
 
         # series
-        _maybe_remove(store, "s")
-        s = Series(np.zeros(12), index=make_index(["date", None, None]))
-        store.append("s", s)
+        _maybe_remove(store, "ser")
+        ser = Series(np.zeros(12), index=make_index(["date", None, None]))
+        store.append("ser", ser)
         xp = Series(np.zeros(12), index=make_index(["date", "level_1", "level_2"]))
-        tm.assert_series_equal(store.select("s"), xp)
+        tm.assert_series_equal(store.select("ser"), xp)
 
         # dup with column
         _maybe_remove(store, "df")

Diff for: pandas/tests/io/pytables/test_read.py

@@ -154,7 +154,7 @@ def test_pytables_native_read(datapath):
         datapath("io", "data", "legacy_hdf/pytables_native.h5"), mode="r"
     ) as store:
         d2 = store["detector/readout"]
-        assert isinstance(d2, DataFrame)
+    assert isinstance(d2, DataFrame)
 
 
 @pytest.mark.skipif(is_platform_windows(), reason="native2 read fails oddly on windows")
@@ -164,7 +164,7 @@ def test_pytables_native2_read(datapath):
     ) as store:
         str(store)
         d1 = store["detector"]
-        assert isinstance(d1, DataFrame)
+    assert isinstance(d1, DataFrame)
 
 
 def test_legacy_table_fixed_format_read_py2(datapath):
@@ -174,28 +174,29 @@ def test_legacy_table_fixed_format_read_py2(datapath):
         datapath("io", "data", "legacy_hdf", "legacy_table_fixed_py2.h5"), mode="r"
     ) as store:
         result = store.select("df")
-        expected = DataFrame(
-            [[1, 2, 3, "D"]],
-            columns=["A", "B", "C", "D"],
-            index=Index(["ABC"], name="INDEX_NAME"),
-        )
-        tm.assert_frame_equal(expected, result)
+    expected = DataFrame(
+        [[1, 2, 3, "D"]],
+        columns=["A", "B", "C", "D"],
+        index=Index(["ABC"], name="INDEX_NAME"),
+    )
+    tm.assert_frame_equal(expected, result)
 
 
 def test_legacy_table_fixed_format_read_datetime_py2(datapath):
     # GH 31750
     # legacy table with fixed format and datetime64 column written in Python 2
+    expected = DataFrame(
+        [[Timestamp("2020-02-06T18:00")]],
+        columns=["A"],
+        index=Index(["date"]),
+        dtype="M8[ns]",
+    )
     with ensure_clean_store(
         datapath("io", "data", "legacy_hdf", "legacy_table_fixed_datetime_py2.h5"),
         mode="r",
     ) as store:
         result = store.select("df")
-        expected = DataFrame(
-            [[Timestamp("2020-02-06T18:00")]],
-            columns=["A"],
-            index=Index(["date"]),
-        )
-        tm.assert_frame_equal(expected, result)
+    tm.assert_frame_equal(expected, result)
 
 
 def test_legacy_table_read_py2(datapath):
@@ -264,7 +265,7 @@ def test_read_hdf_iterator(tmp_path, setup_path):
     with closing(iterator.store):
         assert isinstance(iterator, TableIterator)
         indirect = next(iterator.__iter__())
-        tm.assert_frame_equal(direct, indirect)
+    tm.assert_frame_equal(direct, indirect)
 
 
 def test_read_nokey(tmp_path, setup_path):
@@ -387,7 +388,7 @@ def test_read_py2_hdf_file_in_py3(datapath):
         mode="r",
     ) as store:
         result = store["p"]
-        tm.assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_read_infer_string(tmp_path, setup_path):

Diff for: pandas/tests/io/pytables/test_retain_attributes.py

@@ -1,9 +1,8 @@
 import pytest
 
-from pandas._libs.tslibs import Timestamp
-
 from pandas import (
     DataFrame,
+    DatetimeIndex,
     Series,
     _testing as tm,
     date_range,
@@ -18,11 +17,10 @@
 pytestmark = pytest.mark.single_cpu
 
 
-def test_retain_index_attributes(setup_path):
+def test_retain_index_attributes(setup_path, unit):
     # GH 3499, losing frequency info on index recreation
-    df = DataFrame(
-        {"A": Series(range(3), index=date_range("2000-1-1", periods=3, freq="h"))}
-    )
+    dti = date_range("2000-1-1", periods=3, freq="h", unit=unit)
+    df = DataFrame({"A": Series(range(3), index=dti)})
 
     with ensure_clean_store(setup_path) as store:
         _maybe_remove(store, "data")
@@ -37,37 +35,30 @@ def test_retain_index_attributes(setup_path):
                     getattr(result, idx), attr, None
                 )
 
+        dti2 = date_range("2002-1-1", periods=3, freq="D", unit=unit)
         # try to append a table with a different frequency
         with tm.assert_produces_warning(errors.AttributeConflictWarning):
-            df2 = DataFrame(
-                {
-                    "A": Series(
-                        range(3), index=date_range("2002-1-1", periods=3, freq="D")
-                    )
-                }
-            )
+            df2 = DataFrame({"A": Series(range(3), index=dti2)})
             store.append("data", df2)
 
         assert store.get_storer("data").info["index"]["freq"] is None
 
         # this is ok
         _maybe_remove(store, "df2")
+        dti3 = DatetimeIndex(
+            ["2001-01-01", "2001-01-02", "2002-01-01"], dtype=f"M8[{unit}]"
+        )
         df2 = DataFrame(
             {
                 "A": Series(
                     range(3),
-                    index=[
-                        Timestamp("20010101"),
-                        Timestamp("20010102"),
-                        Timestamp("20020101"),
-                    ],
+                    index=dti3,
                 )
             }
         )
         store.append("df2", df2)
-        df3 = DataFrame(
-            {"A": Series(range(3), index=date_range("2002-1-1", periods=3, freq="D"))}
-        )
+        dti4 = date_range("2002-1-1", periods=3, freq="D", unit=unit)
+        df3 = DataFrame({"A": Series(range(3), index=dti4)})
         store.append("df2", df3)
 
 

Diff for: pandas/tests/io/pytables/test_round_trip.py

@@ -10,6 +10,7 @@
 import pandas as pd
 from pandas import (
     DataFrame,
+    DatetimeIndex,
     Index,
     Series,
     _testing as tm,
@@ -320,7 +321,11 @@ def test_index_types(setup_path):
     ser = Series(values, [1, 5])
     _check_roundtrip(ser, func, path=setup_path)
 
-    ser = Series(values, [datetime.datetime(2012, 1, 1), datetime.datetime(2012, 1, 2)])
+    dti = DatetimeIndex(["2012-01-01", "2012-01-02"], dtype="M8[ns]")
+    ser = Series(values, index=dti)
+    _check_roundtrip(ser, func, path=setup_path)
+
+    ser.index = ser.index.as_unit("s")
     _check_roundtrip(ser, func, path=setup_path)
 
 

Diff for: pandas/tests/io/pytables/test_store.py

@@ -542,17 +542,16 @@ def test_store_index_name(setup_path):
 
 
 @pytest.mark.parametrize("tz", [None, "US/Pacific"])
-@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"])
 @pytest.mark.parametrize("table_format", ["table", "fixed"])
 def test_store_index_name_numpy_str(tmp_path, table_format, setup_path, unit, tz):
     # GH #13492
-    idx = Index(
-        pd.to_datetime([dt.date(2000, 1, 1), dt.date(2000, 1, 2)]),
+    idx = DatetimeIndex(
+        [dt.date(2000, 1, 1), dt.date(2000, 1, 2)],
         name="cols\u05d2",
     ).tz_localize(tz)
     idx1 = (
-        Index(
-            pd.to_datetime([dt.date(2010, 1, 1), dt.date(2010, 1, 2)]),
+        DatetimeIndex(
+            [dt.date(2010, 1, 1), dt.date(2010, 1, 2)],
             name="rows\u05d0",
         )
         .as_unit(unit)

Diff for: pandas/tests/io/pytables/test_time_series.py

@@ -5,6 +5,7 @@
 
 from pandas import (
     DataFrame,
+    DatetimeIndex,
     Series,
     _testing as tm,
 )
@@ -13,10 +14,12 @@
 pytestmark = pytest.mark.single_cpu
 
 
-def test_store_datetime_fractional_secs(setup_path):
+@pytest.mark.parametrize("unit", ["us", "ns"])
+def test_store_datetime_fractional_secs(setup_path, unit):
+    dt = datetime.datetime(2012, 1, 2, 3, 4, 5, 123456)
+    dti = DatetimeIndex([dt], dtype=f"M8[{unit}]")
+    series = Series([0], index=dti)
     with ensure_clean_store(setup_path) as store:
-        dt = datetime.datetime(2012, 1, 2, 3, 4, 5, 123456)
-        series = Series([0], [dt])
         store["a"] = series
         assert store["a"].index[0] == dt
 

Diff for: pandas/tests/io/pytables/test_timezones.py

@@ -148,16 +148,20 @@ def test_append_with_timezones_as_index(setup_path, gettz):
         tm.assert_frame_equal(result, df)
 
 
-def test_roundtrip_tz_aware_index(setup_path):
+def test_roundtrip_tz_aware_index(setup_path, unit):
     # GH 17618
-    time = Timestamp("2000-01-01 01:00:00", tz="US/Eastern")
-    df = DataFrame(data=[0], index=[time])
+    ts = Timestamp("2000-01-01 01:00:00", tz="US/Eastern")
+    dti = DatetimeIndex([ts]).as_unit(unit)
+    df = DataFrame(data=[0], index=dti)
 
     with ensure_clean_store(setup_path) as store:
         store.put("frame", df, format="fixed")
         recons = store["frame"]
         tm.assert_frame_equal(recons, df)
-        assert recons.index[0]._value == 946706400000000000
+
+    value = recons.index[0]._value
+    denom = {"ns": 1, "us": 1000, "ms": 10**6, "s": 10**9}[unit]
+    assert value == 946706400000000000 // denom
 
 
 def test_store_index_name_with_tz(setup_path):
@@ -365,7 +369,7 @@ def test_py2_created_with_datetimez(datapath):
     # Python 3.
     #
     # GH26443
-    index = [Timestamp("2019-01-01T18:00").tz_localize("America/New_York")]
+    index = DatetimeIndex(["2019-01-01T18:00"], dtype="M8[ns, America/New_York]")
     expected = DataFrame({"data": 123}, index=index)
     with ensure_clean_store(
         datapath("io", "data", "legacy_hdf", "gh26443.h5"), mode="r"