diff --git a/src/buddhagpt/corpus.py b/src/buddhagpt/corpus.py index 4c6e878..00f832d 100644 --- a/src/buddhagpt/corpus.py +++ b/src/buddhagpt/corpus.py @@ -23,9 +23,11 @@ def load_bilara_suttas(root: Path) -> Iterator[dict]: segments = json.loads(f.read_text()) uid = f.name.split("_")[0] title = next(iter(segments.values()), uid) + # Derive collection from first directory under sutta/ (e.g., an, dn, kn, mn, sn) + collection = f.relative_to(base).parts[0] yield { "uid": uid, "title": title.strip(), "text": segments_to_text(segments), - "collection": f.parent.name, + "collection": collection, } diff --git a/tests/test_corpus.py b/tests/test_corpus.py index c15f5c7..aa4a2b2 100644 --- a/tests/test_corpus.py +++ b/tests/test_corpus.py @@ -6,7 +6,22 @@ def test_segments_to_text_joins_in_key_order(): segs = {"mn21:1.2": "second.", "mn21:1.1": "First."} assert segments_to_text(segs) == "First. second." +def test_segments_to_text_handles_range_keys(): + """Regression test: range keys like 'an1.142-149:1.0' must not crash. + This would fail with the brief's original int/str-mixing sort key.""" + segs = { + "an1.142-149:1.0": "Range intro.", + "an1.142-149:1.1": "First sentence.", + "an1.140:1.0": "Earlier part.", + } + result = segments_to_text(segs) + # Should not raise TypeError; should contain all texts + assert "Range intro." in result + assert "First sentence." in result + assert "Earlier part." in result + def test_load_bilara_suttas(tmp_path): + """Original test: flat structure.""" d = tmp_path / "translation/en/sujato/sutta/mn" d.mkdir(parents=True) (d / "mn21_translation-en-sujato.json").write_text( @@ -15,3 +30,17 @@ def test_load_bilara_suttas(tmp_path): suttas = list(load_bilara_suttas(tmp_path)) assert suttas[0]["uid"] == "mn21" assert "So I have heard." in suttas[0]["text"] + assert suttas[0]["collection"] == "mn" + +def test_load_bilara_suttas_nested_layout(tmp_path): + """Regression test: nested structure (an/an1/, sn/sn35/, etc.) + verifies collection is extracted from top-level nikaya, not immediate parent.""" + d = tmp_path / "translation/en/sujato/sutta/an/an1" + d.mkdir(parents=True) + (d / "an1.1-10_translation-en-sujato.json").write_text( + json.dumps({"an1.1-10:0.1": "AN 1.1-10", "an1.1-10:1.1": "Content."}) + ) + suttas = list(load_bilara_suttas(tmp_path)) + assert suttas[0]["uid"] == "an1.1-10" + assert suttas[0]["collection"] == "an", f"Expected 'an', got '{suttas[0]['collection']}'" + assert "Content." in suttas[0]["text"]