Repository: apache/arrow
File: python/pyarrow/tests/test_pandas.py (Python)

Code before:
    # with no loss of information if the timestamp_as_object option is True.
    df = make_df_with_timestamps()
    table = pa.Table.from_pandas(df)
    df2 = table.to_pandas(timestamp_as_object=True)
    tm.assert_frame_equal(df, df2)


@pytest.mark.parametrize("resolution", ["s", "ms", "us"])
@pytest.mark.parametrize("tz", [None, "America/New_York"])
# One datetime outside nanosecond range, one inside nanosecond range:
@pytest.mark.parametrize("dt", [datetime(1553, 1, 1), datetime(2020, 1, 1)])
def test_timestamp_as_object_non_nanosecond(resolution, tz, dt):
    # Timestamps can be converted Arrow and reloaded into Pandas with no loss
    # of information if the timestamp_as_object option is True.
    arr = pa.array([dt], type=pa.timestamp(resolution, tz=tz))
    table = pa.table({'a': arr})

    for result in [
        arr.to_pandas(timestamp_as_object=True),
        table.to_pandas(timestamp_as_object=True)['a']
    ]:
        assert result.dtype == object
        assert isinstance(result[0], datetime)
        if tz:
            assert result[0].tzinfo is not None
            expected = dt.replace(tzinfo=timezone.utc).astimezone(result[0].tzinfo)
        else:
            assert result[0].tzinfo is None
            expected = dt
        assert result[0] == expected


def test_timestamp_as_object_fixed_offset():
    # ARROW-16547 to_pandas with timestamp_as_object=True and FixedOffset
    pytz = pytest.importorskip("pytz")
    import datetime
    timezone = pytz.FixedOffset(120)
    dt = timezone.localize(datetime.datetime(2022, 5, 12, 16, 57))

    table = pa.table({"timestamp_col": pa.array([dt])})
    result = table.to_pandas(timestamp_as_object=True)
    assert pa.table(result) == table


@pytest.mark.processes
def test_threaded_pandas_import():
    invoke_script("pandas_threaded_import.py")


def test_does_not_mutate_timedelta_dtype():
    expected = np.dtype('m8')

Review comment (suggestion):
This is totally legit, but we might want to use `tz` just to make it cleaner:


```suggestion
            expected = dt.replace(tzinfo=timezone.utc).astimezone(ZoneInfo(tz))
```

and we would need `from zoneinfo import ZoneInfo` on top of the file.

Code after:
    # Out of range timestamps can be converted Arrow and reloaded into Pandas
    # with no loss of information if the timestamp_as_object option is True.
    df = make_df_with_timestamps()
    table = pa.Table.from_pandas(df)
    df2 = table.to_pandas(timestamp_as_object=True)
    tm.assert_frame_equal(df, df2)


@pytest.mark.parametrize("resolution", ["s", "ms", "us"])
@pytest.mark.parametrize("tz", [None, "America/New_York"])
# One datetime outside nanosecond range, one inside nanosecond range:
@pytest.mark.parametrize("dt", [datetime(1553, 1, 1), datetime(2020, 1, 1)])
def test_timestamp_as_object_non_nanosecond(resolution, tz, dt):
    # Timestamps can be converted Arrow and reloaded into Pandas with no loss
    # of information if the timestamp_as_object option is True.
    arr = pa.array([dt], type=pa.timestamp(resolution, tz=tz))
    table = pa.table({'a': arr})

    for result in [
        arr.to_pandas(timestamp_as_object=True),
        table.to_pandas(timestamp_as_object=True)['a']
    ]:
        assert result.dtype == object
        assert isinstance(result[0], datetime)
        if tz:
            assert result[0].tzinfo is not None
            expected = dt.replace(tzinfo=timezone.utc).astimezone(ZoneInfo(tz))
        else:
            assert result[0].tzinfo is None
            expected = dt
        assert result[0] == expected


def test_timestamp_as_object_fixed_offset():
    # ARROW-16547 to_pandas with timestamp_as_object=True and FixedOffset
    pytz = pytest.importorskip("pytz")
    import datetime
    timezone = pytz.FixedOffset(120)
    dt = timezone.localize(datetime.datetime(2022, 5, 12, 16, 57))

    table = pa.table({"timestamp_col": pa.array([dt])})
    result = table.to_pandas(timestamp_as_object=True)
    assert pa.table(result) == table


@pytest.mark.processes
def test_threaded_pandas_import():
    invoke_script("pandas_threaded_import.py")


def test_does_not_mutate_timedelta_dtype():