"""
Testing utilities for Polars DataFrame snapshots.
This module provides a custom snapshot testing implementation that stores
Polars DataFrame snapshots as Parquet files and uses polars.testing.assert_frame_equal
for comparison with clear diff output.
"""
from pathlib import Path
from typing import TYPE_CHECKING
# For type checking, import polars types
if TYPE_CHECKING:
import polars
# Try to import polars, but allow the module to be imported even if polars is not available
try:
import polars as pl
from polars.testing import assert_frame_equal
_POLARS_AVAILABLE = True
except ImportError:
pl = None
assert_frame_equal = None
_POLARS_AVAILABLE = False
[docs]
class DataframeSnapshot:
"""
Class to manage Polars DataFrame snapshots stored as Parquet files.
This class provides snapshot testing functionality for Polars DataFrames by storing
them as Parquet files and comparing them using polars.testing.assert_frame_equal.
The parent directory for the snapshot file is automatically created if it doesn't exist.
Args:
path (Path | str): Complete file path to the snapshot file
is_snapshot_update (bool, optional): If True, snapshots will be updated instead
of compared. Defaults to False.
Usage:
Typically used with the df_snapshot pytest fixture::
def test_my_dataframe(df_snapshot: DataframeSnapshot):
df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
df_snapshot.assert_frame_equal(df)
Or create directly::
snapshot = DataframeSnapshot("example-output/unit/test_df_snapshot.parquet")
df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
snapshot.assert_frame_equal(df)
To update snapshots, run the test with --df-snapshot-update
Raises:
ImportError: If polars is not installed when trying to use the class
"""
def __init__(self, path: Path | str, is_snapshot_update: bool = False):
"""
Initialize DataFrameSnapshot instance.
Args:
path: File path where the snapshot will be stored
is_snapshot_update: If True, snapshots will be updated instead of compared
Raises:
ImportError: If polars is not installed
"""
if not _POLARS_AVAILABLE:
raise ImportError(
"polars is required to use DataframeSnapshot. "
"Install it with: pip install 'ts-lib-pytest[polars]'"
)
self.snapshot_path = Path(path)
# Ensure the parent directory exists
self.snapshot_path.parent.mkdir(parents=True, exist_ok=True)
self.is_snapshot_update = is_snapshot_update
[docs]
def read_snapshot(self) -> "polars.DataFrame":
"""Read the snapshot from file."""
df = pl.read_parquet(self.snapshot_path) # type: ignore
return df
def _save_snapshot(self, df: "polars.DataFrame") -> None:
"""Save the snapshot to file."""
df.write_parquet(self.snapshot_path)
[docs]
def assert_frame_equal(
self,
df_right: "polars.DataFrame",
*,
check_row_order: bool = True,
check_column_order: bool = True,
check_dtypes: bool = True,
check_exact: bool = False,
rel_tol: float = 1e-5,
abs_tol: float = 1e-8,
categorical_as_str: bool = False,
) -> None:
"""
Compare DataFrame to snapshot or update the snapshot.
If this is the first run and no snapshot exists, the snapshot will be created
and the test will fail to alert the developer. If is_snapshot_update is True,
the snapshot will be updated instead of compared.
Args:
df_right: The DataFrame to compare against the snapshot
check_row_order: Whether to check that rows are in the same order
check_column_order: Whether to check that columns are in the same order
check_dtypes: Whether to check that data types match
check_exact: Whether to check exact equality for floating point values
rel_tol: Relative tolerance for floating point comparisons
abs_tol: Absolute tolerance for floating point comparisons
categorical_as_str: Whether to compare categorical columns as strings
Raises:
FileNotFoundError: If snapshot doesn't exist and is created for the first time
AssertionError: If the DataFrame doesn't match the snapshot
"""
if self.is_snapshot_update:
self._save_snapshot(df_right)
return
try:
snapshot_df = self.read_snapshot()
except FileNotFoundError:
self._save_snapshot(df_right)
# Fail the test after creating the snapshot to alert the developer
raise FileNotFoundError(
f"Snapshot was missing and has been created at {self.snapshot_path}. "
f"This test failure alerts you in case the snapshot was created inadvertently. "
f"Run the test again to verify the function produces consistent results."
)
try:
assert_frame_equal( # type: ignore
snapshot_df,
df_right,
check_row_order=check_row_order,
check_column_order=check_column_order,
check_dtypes=check_dtypes,
check_exact=check_exact,
rel_tol=rel_tol,
abs_tol=abs_tol,
categorical_as_str=categorical_as_str,
)
except AssertionError as e:
# Add a message to the assertion error to help developers update the snapshot
raise AssertionError(
f"Snapshot test failed. Run pytest with --df-snapshot-update to update the snapshot.\n"
f"{str(e)}"
)