Coverage for polypandas/testing.py: 88%

1"""Testing utilities for pandas DataFrames."""

2

3 from typing import Any, Dict, List, Optional

4

5 from polypandas.exceptions import PolypandasError

6 from polypandas.protocols import is_pandas_available

7

8

9 class DataFrameComparisonError(PolypandasError):

10 """Raised when DataFrame comparison fails."""

11

12 pass

13

14

15 def assert_dataframe_equal(

16 df1: Any,

17 df2: Any,

18 check_order: bool = False,

19 rtol: float = 1e-5,

20 atol: float = 1e-8,

21 check_dtypes: bool = True,

22 check_column_order: bool = False,

23) -> None:

24 """Assert that two pandas DataFrames are equal.

25

26 Args:

27 df1: First DataFrame to compare.

28 df2: Second DataFrame to compare.

29 check_order: If True, row order must match. If False, DataFrames are sorted.

30 rtol: Relative tolerance for floating point comparisons.

31 atol: Absolute tolerance for floating point comparisons.

32 check_dtypes: If True, check that dtypes match.

33 check_column_order: If True, column order must match.

34

35 Raises:

36 DataFrameComparisonError: If DataFrames are not equal.

37 """

38 if not is_pandas_available():

39 raise DataFrameComparisonError(

40 "pandas is required for DataFrame comparison. Install it with: pip install pandas"

41 )

42

43 import pandas as pd

44

45 if not isinstance(df1, pd.DataFrame) or not isinstance(df2, pd.DataFrame):

46 raise DataFrameComparisonError("Both arguments must be pandas DataFrames")

47

48 if len(df1) != len(df2):

49 raise DataFrameComparisonError(

50 f"DataFrame row counts don't match: {len(df1)} != {len(df2)}"

51 )

52

53 if set(df1.columns) != set(df2.columns):

54 raise DataFrameComparisonError(

55 f"DataFrame columns don't match: {set(df1.columns)} != {set(df2.columns)}"

56 )

57

58 if check_column_order and list(df1.columns) != list(df2.columns):

59 raise DataFrameComparisonError("Column order does not match")

60

61 if not check_column_order:

62 df2 = df2[df1.columns]

63

64 if not check_order and len(df1) > 0:

65 df1 = df1.sort_values(by=list(df1.columns)).reset_index(drop=True)

66 df2 = df2.sort_values(by=list(df2.columns)).reset_index(drop=True)

67

68 try:

69 pd.testing.assert_frame_equal(

70 df1,

71 df2,

72 check_dtype=check_dtypes,

73 rtol=rtol,

74 atol=atol,

75 )

76 except AssertionError as e:

77 raise DataFrameComparisonError(str(e)) from e

78

79

80 def assert_schema_equal(

81 df1: Any,

82 df2: Any,

83 check_order: bool = False,

84) -> None:

85 """Assert that two DataFrames have the same dtypes (schema). Alias for assert_dtypes_equal."""

86 assert_dtypes_equal(df1, df2, check_order=check_order)

87

88

89 def assert_dtypes_equal(

90 df1: Any,

91 df2: Any,

92 check_order: bool = False,

93) -> None:

94 """Assert that two DataFrames have the same dtypes.

95

96 Args:

97 df1: First DataFrame.

98 df2: Second DataFrame.

99 check_order: If True, column order must match.

100

101 Raises:

102 DataFrameComparisonError: If dtypes are not equal.

103 """

104 if not is_pandas_available():

105 raise DataFrameComparisonError("pandas is required")

106

107 dtypes1 = df1.dtypes

108 dtypes2 = df2.dtypes

109

110 if set(dtypes1.index) != set(dtypes2.index):

111 raise DataFrameComparisonError("DataFrames have different columns")

112

113 if not check_order:

114 dtypes2 = dtypes2[dtypes1.index]

115

116 for col in dtypes1.index:

117 if dtypes1[col] != dtypes2[col]:

118 raise DataFrameComparisonError(

119 f"Column '{col}' has different dtypes: {dtypes1[col]} != {dtypes2[col]}"

120 )

121

122

123 def assert_approx_count(df: Any, expected_count: int, tolerance: float = 0.1) -> None:

124 """Assert that DataFrame row count is approximately equal to expected."""

125 if not is_pandas_available():

126 raise DataFrameComparisonError("pandas is required")

127

128 actual_count = len(df)

129 min_count = int(expected_count * (1 - tolerance))

130 max_count = int(expected_count * (1 + tolerance))

131

132 if not (min_count <= actual_count <= max_count):

133 raise DataFrameComparisonError(

134 f"DataFrame count {actual_count} is not within {tolerance * 100:.1f}% "

135 f"of expected {expected_count}. Expected range: [{min_count}, {max_count}]"

136 )

137

138

139 def get_column_stats(df: Any, column: str) -> Dict[str, Any]:

140 """Get basic statistics for a column."""

141 if not is_pandas_available():

142 raise DataFrameComparisonError("pandas is required")

143

144 import pandas as pd

145

146 s = df[column]

147 stats: Dict[str, Any] = {

148 "count": int(s.count()),

149 "null_count": int(s.isna().sum()),

150 "distinct_count": int(s.nunique()),

151 }

152

153 if pd.api.types.is_numeric_dtype(s):

154 stats["min"] = float(s.min()) if s.count() else None

155 stats["max"] = float(s.max()) if s.count() else None

156 stats["mean"] = float(s.mean()) if s.count() else None

157 stats["std"] = float(s.std()) if s.count() else None

158

159 return stats

160

161

162 def assert_column_exists(df: Any, *columns: str) -> None:

163 """Assert that specified columns exist in DataFrame."""

164 df_columns = set(df.columns)

165 missing = [c for c in columns if c not in df_columns]

166

167 if missing:

168 raise DataFrameComparisonError(

169 f"Columns missing from DataFrame: {missing}. Available columns: {sorted(df_columns)}"

170 )

171

172

173 def assert_no_duplicates(df: Any, columns: Optional[List[str]] = None) -> None:

174 """Assert that DataFrame has no duplicate rows."""

175 if not is_pandas_available():

176 raise DataFrameComparisonError("pandas is required")

177

178 total_count = len(df)

179 if columns is None:

180 unique_count = len(df.drop_duplicates())

181 else:

182 unique_count = len(df.drop_duplicates(subset=columns))

183

184 if unique_count != total_count:

185 raise DataFrameComparisonError(

186 f"DataFrame contains {total_count - unique_count} duplicate row(s). "

187 f"Total rows: {total_count}, Unique rows: {unique_count}"

188 )