String to datetime conversion with custom format about polars HOT 5 OPEN

from datetime import datetime

import pandas as pd
import polars as pl
import pytest

dt_string = "2024-06-03 20:02:48.6800000"


dt_formats = [
    "%Y-%m-%d %H:%M:%S%.f",
    "%Y-%m-%d %H:%M:%S%.6f",
    "%Y-%m-%d %H:%M:%S.%f",
    "%Y-%m-%d %H:%M:%S.%6f",
]
dt_formats = dt_formats + [f + "0" for f in dt_formats]


@pytest.fixture
def df():
    return pl.DataFrame({"dt": dt_string})


@pytest.mark.parametrize("format", dt_formats)
def test_polars(df, format, capsys):
    with capsys.disabled():
        print(format)
        print(df)

    t_ref = pl.Series(name="dt", values=[datetime(2024, 6, 3, 20, 2, 48, 680000)])
    df2 = df.with_columns(dt=pl.col("dt").str.to_datetime(format, time_unit="ns"))

    with capsys.disabled():
        print(df2)
    assert (df2["dt"] == t_ref).all()


@pytest.mark.parametrize("format", dt_formats)
def test_pandas(df, format, capsys):
    with capsys.disabled():
        print(format)
        print(df)

    t_ref = pd.Series(name="dt", data=[datetime(2024, 6, 3, 20, 2, 48, 680000)])
    t = pd.to_datetime(
        df["dt"].to_pandas(use_pyarrow_extension_array=True), format=format
    )

    with capsys.disabled():
        print(t)

    assert (t == t_ref).all()


@pytest.mark.parametrize("format", dt_formats)
def test_datetime(df, format, capsys):
    dt_ref = datetime(2024, 6, 3, 20, 2, 48, 680000)
    dt = datetime.strptime(dt_string, format)

    assert dt == dt_ref