Fix #14: k_fold covers every row in exactly one test fold

mess_mood/evaluate.py

@@ -16,11 +16,14 @@ def k_fold(rows, k=5, seed=42):
     """Yields (train, test) k times. Every row is in exactly one test set."""
     rows = rows[:]
     random.Random(seed).shuffle(rows)
-    fold_size = len(rows) // k
+    fold_size, extra = divmod(len(rows), k)
+    start = 0
     for i in range(k):
-        test = rows[i * fold_size:(i + 1) * fold_size]
-        train = rows[:i * fold_size] + rows[(i + 1) * fold_size:]
+        end = start + fold_size + (i < extra)
+        test = rows[start:end]
+        train = rows[:start] + rows[end:]
         yield train, test
+        start = end
 
 
 def fit(rows):

tests/test_mess_mood.py

@@ -68,6 +68,29 @@ def test_k_fold_test_sets_dont_overlap():
     assert all(not (a & b) for i, a in enumerate(tests) for b in tests[i + 1:])
 
 
+def test_k_fold_every_row_in_exactly_one_test_set():
+    from collections import Counter
+
+    for n in (23, 24, 21, 7):
+        rows = [(str(i), "x") for i in range(n)]
+        folds = list(k_fold(rows, k=5))
+        assert len(folds) == 5
+        assert Counter(r for _, test in folds for r in test) == Counter(rows)
+        for train, test in folds:
+            assert sorted(train + test) == sorted(rows)
+            assert not set(train) & set(test)
+
+
+def test_k_fold_covers_every_shipped_review():
+    from collections import Counter
+
+    rows = load_reviews()
+    folds = list(k_fold(rows, k=5))
+    assert Counter(r for _, test in folds for r in test) == Counter(rows)
+    sizes = [len(test) for _, test in folds]
+    assert max(sizes) - min(sizes) <= 1
+
+
 def test_model_save_and_load(tmp_path):
     from mess_mood.model import NaiveBayes