experimental assert_frame_equal with snapshot support

inline_snapshot/_inline_snapshot.py

@@ -237,7 +237,7 @@ def __eq__(self, other):
         if self._new_value is undefined:
             self._new_value = other
 
-        return self._visible_value() == other
+        return other == self._visible_value()
 
     def _new_code(self):
         return self._value_to_code(self._new_value)
@@ -355,7 +355,7 @@ def check(old_value, old_node, new_value):
             # generic fallback
             new_token = value_to_token(new_value)
 
-            if not old_value == new_value:
+            if not new_value == old_value:
                 flag = "fix"
             elif (
                 self._ast_node is not None

inline_snapshot/_pandas.py

@@ -0,0 +1,57 @@
+from functools import wraps
+from typing import Optional
+
+
+def make_assert_equal(data_type, assert_equal, repr_function):
+
+    class Wrapper:
+        def __init__(self, df, cmp):
+            self.df = df
+            self.cmp = cmp
+
+        def __repr__(self):
+            return f"{data_type.__name__}({repr(repr_function(self.df))})"
+
+        def __eq__(self, other):
+            if not isinstance(other, data_type):
+                return NotImplemented
+            return self.cmp(self.df, other)
+
+    @wraps(assert_equal)
+    def result(df, df_snapshot, *args, **kargs):
+        error: Optional[AssertionError] = None
+
+        def cmp(a, b):
+            nonlocal error
+            try:
+                assert_equal(a, b, *args, **kargs)
+            except AssertionError as e:
+                error = e
+                return False
+            return True
+
+        if not Wrapper(df, cmp) == df_snapshot:
+            assert error is not None
+            raise error
+
+    return result
+
+
+try:
+    import pandas
+except:
+    pass
+else:
+    from pandas.testing import assert_frame_equal
+    from pandas.testing import assert_index_equal
+    from pandas.testing import assert_series_equal
+
+    assert_frame_equal = make_assert_equal(
+        pandas.DataFrame, assert_frame_equal, lambda df: df.to_dict("records")
+    )
+    assert_series_equal = make_assert_equal(
+        pandas.Series, assert_series_equal, lambda df: df.to_dict()
+    )
+    assert_index_equal = make_assert_equal(
+        pandas.Index, assert_index_equal, lambda df: df.to_list()
+    )

pyproject.toml

@@ -56,6 +56,7 @@ asttokens = ">=2.0.5"
 black = ">=23.3.0"
 click = ">=8.1.4"
 executing = ">=2.0.0"
+pandas = {version = "^2.2.2", python = ">=3.9"}
 python = ">=3.8"
 rich = ">=13.7.1"
 toml = ">=0.10.2"

tests/test_pandas.py

@@ -0,0 +1,42 @@
+import sys
+
+import pytest
+
+if sys.version_info >= (3, 9):
+    from pandas import DataFrame
+    from pandas import Index
+    from pandas import Series
+
+    from inline_snapshot import snapshot
+    from inline_snapshot._pandas import assert_frame_equal
+    from inline_snapshot._pandas import assert_index_equal
+    from inline_snapshot._pandas import assert_series_equal
+
+
+@pytest.mark.skipif(sys.version_info < (3, 9), reason="no pandas for 3.9")
+def test_df():
+    df = DataFrame({"col0": [1, 2], "col1": [1, 5j], "col3": ["a", "b"]})
+
+    # the second argument can be a snapshot
+    assert_frame_equal(
+        df,
+        snapshot(
+            DataFrame(
+                [
+                    {"col0": 1, "col1": (1 + 0j), "col3": "a"},
+                    {"col0": 2, "col1": 5j, "col3": "b"},
+                ]
+            )
+        ),
+    )
+
+    # and can also be used without a snapshot
+    assert_frame_equal(df, df)
+
+    # for Index
+    index = Index(range(5))
+    assert_index_equal(index, snapshot(Index([0, 1, 2, 3, 4])))
+
+    # for Series
+    index = Series({1: 8, 5: 4})
+    assert_series_equal(index, snapshot(Series({1: 8, 5: 4})))