Coverage for src / inline_snapshot_pandas / __init__.py: 100%

1from functools import wraps 

2from typing import Optional 

3 

4import pandas.testing 

5from inline_snapshot import customize_repr 

6from inline_snapshot import snapshot 

7from inline_snapshot._snapshot.generic_value import GenericValue 

8 

9 

10__all__ = ( 

11 "setup", 

12 "assert_frame_equal", 

13 "assert_series_equal", 

14 "assert_index_equal", 

15 "snapshot", 

16) 

17 

18 

19def make_assert_equal(data_type, assert_equal, repr_function): 

20 

21 class Wrapper: 

22 def __init__(self, df, cmp): 

23 self.df = df 

24 self.cmp = cmp 

25 

26 def __repr__(self): 

27 return f"{data_type.__name__}({repr_function(self.df)!r})" 

28 

29 def __eq__(self, other): 

30 if isinstance(other, data_type): 

31 return self.cmp(self.df, other) 

32 if isinstance(other, Wrapper) and isinstance(other.df, data_type): 

33 return self.cmp(self.df, other.df) 

34 return NotImplemented 

35 

36 original = data_type.__eq__ 

37 

38 def new_eq(a, b): 

39 if isinstance(b, (GenericValue, Wrapper)): 

40 return NotImplemented 

41 return original(a, b) 

42 

43 data_type.__eq__ = new_eq 

44 

45 @wraps(assert_equal) 

46 def result(df, df_snapshot, *args, **kargs): 

47 error: Optional[AssertionError] = None 

48 

49 def cmp(a, b): 

50 nonlocal error 

51 try: 

52 assert_equal(a, b, *args, **kargs) 

53 except AssertionError as e: 

54 error = e 

55 return False 

56 return True 

57 

58 if not Wrapper(df, cmp) == df_snapshot: 

59 assert error is not None 

60 raise error 

61 

62 return result 

63 

64 

65assert_frame_equal = make_assert_equal( 

66 pandas.DataFrame, 

67 pandas.testing.assert_frame_equal, 

68 lambda df: df.to_dict("records"), 

69) 

70assert_series_equal = make_assert_equal( 

71 pandas.Series, pandas.testing.assert_series_equal, lambda df: df.to_dict() 

72) 

73assert_index_equal = make_assert_equal( 

74 pandas.Index, pandas.testing.assert_index_equal, lambda df: df.to_list() 

75) 

76 

77 

78def setup(): 

79 pandas.testing.assert_frame_equal = assert_frame_equal 

80 pandas.testing.assert_series_equal = assert_series_equal 

81 pandas.testing.assert_index_equal = assert_index_equal