pytest-regtest 2.2.0a2__py2.py3-none-any.whl → 2.3.0__py2.py3-none-any.whl

pytest_regtest/__init__.py

@@ -2,13 +2,24 @@ from importlib.metadata import version as _version
 
 import pytest
 
-from . import register_third_party_handlers  # noqa: F401
-from .pytest_regtest import PytestRegtestPlugin  # noqa: F401
-from .pytest_regtest import RegtestStream  # noqa: F401
 from .pytest_regtest import clear_converters  # noqa: F401
 from .pytest_regtest import patch_terminal_size  # noqa: F401
 from .pytest_regtest import register_converter_post  # noqa: F401
 from .pytest_regtest import register_converter_pre  # noqa: F401
+from .pytest_regtest import (
+    PytestRegtestCommonHooks,
+    PytestRegtestPlugin,
+    RegtestStream,
+    Snapshot,
+    SnapshotPlugin,
+)
+from .register_third_party_handlers import (
+    register_numpy_handler,
+    register_pandas_handler,
+    register_polars_handler,
+)
+
+from .snapshot_handler import register_python_object_handler
 
 __version__ = _version(__package__)
 
@@ -43,12 +54,17 @@ def pytest_addoption(parser):
         "--regtest-disable-stdconv",
         action="store_true",
         default=False,
-        help="do not apply standard output converters to clean up indeterministic output",
+        help=(
+            "do not apply standard output converters to clean up indeterministic output"
+        ),
     )
 
 
 def pytest_configure(config):
-    config.pluginmanager.register(PytestRegtestPlugin())
+    common = PytestRegtestCommonHooks()
+    config.pluginmanager.register(common)
+    config.pluginmanager.register(PytestRegtestPlugin(common))
+    config.pluginmanager.register(SnapshotPlugin(common))
 
 
 @pytest.fixture
@@ -56,7 +72,9 @@ def regtest(request):
     yield RegtestStream(request)
 
 
-snapshot = regtest
+@pytest.fixture
+def snapshot(request):
+    yield Snapshot(request)
 
 
 @pytest.fixture
@@ -65,3 +83,26 @@ def regtest_all(regtest):
 
 
 snapshot_all_output = regtest_all
+
+register_python_object_handler()
+
+try:
+    import pandas  # noqa: F401
+
+    register_pandas_handler()
+except ImportError:
+    pass
+
+try:
+    import numpy  # noqa: F401
+
+    register_numpy_handler()
+except ImportError:
+    pass
+
+try:
+    import polars  # noqa: F401
+
+    register_polars_handler()
+except ImportError:
+    pass

pytest_regtest/numpy_handler.py

@@ -11,12 +11,19 @@ from .utils import highlight_mismatches
 
 class NumpyHandler(BaseSnapshotHandler):
     def __init__(self, handler_options, pytest_config, tw):
-        self.print_options = np.get_printoptions() | handler_options.get(
-            "print_options", {}
-        )
         self.atol = handler_options.get("atol", 0.0)
         self.rtol = handler_options.get("rtol", 0.0)
         self.equal_nan = handler_options.get("equal_nan", True)
+        if handler_options.get("print_options"):
+            warnings.warn(
+                "please use the numpy.printoptions context manager instead of"
+                " the print_options argument.",
+                DeprecationWarning,
+            )
+
+        self.print_options = np.get_printoptions() | handler_options.get(
+            "print_options", {}
+        )
 
     def _filename(self, folder):
         return os.path.join(folder, "arrays.npy")
@@ -148,6 +155,7 @@ class NumpyHandler(BaseSnapshotHandler):
         self, current_obj, current_as_text, recorded_obj, recorded_as_text, has_markup
     ):
         sub_diff = []
+
         for i, (l1, l2, r1, r2) in enumerate(
             zip(current_as_text, recorded_as_text, current_obj, recorded_obj)
         ):
@@ -193,4 +201,16 @@ class NumpyHandler(BaseSnapshotHandler):
                 sub_diff.append(f"row {i:3d}: {l1}")
                 sub_diff.append(f"         {l2}")
 
+        missing = len(current_as_text) - len(recorded_as_text)
+        if missing > 0:
+            for i, row in enumerate(current_as_text[-missing:], len(recorded_as_text)):
+                # remove duplicate brackets
+                row = row.rstrip("]") + "]"
+                sub_diff.append(f"row {i:3d}: -{row.lstrip()}")
+        if missing < 0:
+            for i, row in enumerate(recorded_as_text[missing:], len(current_as_text)):
+                # remove duplicate brackets
+                row = row.rstrip("]") + "]"
+                sub_diff.append(f"row {i:3d}: +{row.lstrip()}")
+
         return sub_diff

pytest_regtest/pandas_handler.py

@@ -1,6 +1,7 @@
 import difflib
 import io
 import os.path
+import warnings
 
 import numpy as np
 import pandas as pd
@@ -10,6 +11,13 @@ from .snapshot_handler import BaseSnapshotHandler
 
 class DataFrameHandler(BaseSnapshotHandler):
     def __init__(self, handler_options, pytest_config, tw):
+        if handler_options.get("display_options"):
+            warnings.warn(
+                "please use the 'pandas.option_context' context manager instead of"
+                " the display_options argument.",
+                DeprecationWarning,
+            )
+
         # default contains a few nested dicts and we flatten those, e.g.
         # { "html": {"border": 1} } -> { "html.border": 1 }
         default = list(pd.options.display.d.items())

pytest_regtest/polars_handler.py

@@ -0,0 +1,114 @@
+import difflib
+import io
+import os
+from typing import Any
+
+import polars as pl
+from polars.testing import assert_frame_equal
+
+from .snapshot_handler import BaseSnapshotHandler
+
+
+class PolarsHandler(BaseSnapshotHandler):
+    """
+    PolarsHandler is a class for handling Polars DataFrame snapshots in pytest-regtest.
+    """
+
+    def __init__(self, handler_options: dict[str, Any], pytest_config, tw):
+        self.atol = handler_options.get("atol", 0.0)
+        self.rtol = handler_options.get("rtol", 0.0)
+        self.display_options = handler_options.get("display_options", None)
+
+    def _filename(self, folder: str | os.PathLike[Any]) -> str:
+        return os.path.join(folder, "polars.parquet")
+
+    def save(self, folder: str | os.PathLike[Any], obj: pl.DataFrame):
+        obj.write_parquet(self._filename(folder))
+
+    def load(self, folder: str | os.PathLike[Any]) -> pl.DataFrame:
+        return pl.read_parquet(self._filename(folder))
+
+    def show(self, obj: pl.DataFrame) -> list[str]:
+        stream = io.StringIO()
+        if self.display_options:
+            with pl.Config(**self.display_options):
+                stream.write(str(obj))
+        else:
+            stream.write(str(obj))
+        return stream.getvalue().splitlines()
+
+    def compare(self, current_obj: pl.DataFrame, recorded_obj: pl.DataFrame) -> bool:
+        try:
+            assert_frame_equal(
+                current_obj, recorded_obj, atol=self.atol, rtol=self.rtol
+            )
+            return True
+        except AssertionError:
+            return False
+
+    @staticmethod
+    def create_schema_info(df: pl.DataFrame) -> list[str]:
+        """
+        Generate a summary of the schema information for a given Polars DataFrame.
+
+        Parameters:
+            df (pl.DataFrame): The Polars DataFrame for which to generate schema information.
+
+        Returns:
+            list[str]: A list of strings representing the schema information, including
+                   the total number of columns, column names, non-null counts, and data types.
+        """
+        schema = df.schema
+        schema_string_repr = [
+            "Data columns (total {} columns):".format(len(schema)),
+            " #   Column  Non-Null Count  Dtype  ",
+            "---  ------  --------------  -----  ",
+        ]
+        for i, (column, dtype) in enumerate(schema.items()):
+            total_count = df.height
+            null_count = df[column].null_count()
+            non_null_count = total_count - null_count
+            dtype_str = str(dtype)
+            schema_string_repr.append(
+                f" {i}   {column}       {non_null_count} non-null      {dtype_str}"
+            )
+        return schema_string_repr
+
+    def show_differences(
+        self, current_obj: pl.DataFrame, recorded_obj: pl.DataFrame, has_markup: bool
+    ) -> list[str]:
+        lines = []
+
+        current_schema = self.create_schema_info(current_obj)
+        recorded_schema = self.create_schema_info(recorded_obj)
+
+        info_diff = list(
+            difflib.unified_diff(
+                current_schema,
+                recorded_schema,
+                "current",
+                "expected",
+                lineterm="",
+            )
+        )
+        lines.extend(info_diff)
+        recorded_as_text = self.show(recorded_obj)
+        current_as_text = self.show(current_obj)
+
+        diffs = list(
+            difflib.unified_diff(
+                current_as_text,
+                recorded_as_text,
+                "current",
+                "expected",
+                lineterm="",
+            )
+        )
+
+        lines.append("")
+        if diffs:
+            lines.extend(diffs)
+        else:
+            lines.append("diff is empty, you may want to change the print options")
+
+        return lines