diff --git a/tools/src/pop_fem_audit_tools/commands/build_db.py b/tools/src/pop_fem_audit_tools/commands/build_db.py index 7beea58..c8c695e 100644 --- a/tools/src/pop_fem_audit_tools/commands/build_db.py +++ b/tools/src/pop_fem_audit_tools/commands/build_db.py @@ -7,59 +7,18 @@ Rebuilds the working store from scratch out of the committed inputs: the year-end chart CSV and the output directory for the review CSV files, given as the two positional command-line -arguments, and the optional inputs, each given as an -option: the lyrics cache directory, the Wikidata artist -snapshot CSV, the settled coding table CSV, and the gender -correction table CSV. An omitted option -leaves its layer unloaded; a given option whose path does -not exist fails the build. Missing tables +arguments, and the optional capture, coding, group, +gender-correction, pattern, and annotation layers, each given as +an option. An omitted option leaves its layer unloaded; a given +option whose path does not exist fails the build. Missing tables are created on a fresh store; existing tables are never altered, as the schema lifecycle belongs to the migrations. Every rebuild deletes all the rows, loads the data, and validates it in one transaction, committed only after the data passes the validation invariants; a failed build leaves the previous store contents -intact. - -The rebuild is deterministic: the builder assigns the song and -artist IDs itself, as 1, 2, 3, ... in the first-occurrence file -order, so the IDs are reproducible across rebuilds on every -database engine, given the frozen input file. - -A song is identified by its raw title together with its artist -credit, the credit canonicalized through -``SongImporter.CANONICAL_ARTIST_CREDITS``; a credit listed there -collapses onto the same song as its canonical form, and the -stored artist credit is always the canonical form. Artist -deduplication is by the identity key resolved from the parsed -artist name (see `ArtistImporter.resolve_artist_identity`): the -case-folded name, or, when that case-folded name is listed in -``ArtistImporter.CANONICAL_ARTIST_NAMES``, the case-folded -canonical spelling, so letter-case variants and alternate -spellings mapped to the same canonical name all collapse onto a -single artist row. The stored artist name is the first-seen -spelling, except for the names listed in -``ArtistImporter.CANONICAL_ARTIST_NAMES``, which always store the -canonical spelling regardless of which variant is seen first. - -Once the artists carry their captured attributes, every song -takes a performer gender derived from the genders of the -performing artists credited on it, primary and featured alike. A -credited artist without an artist type -- a label, a brand, or a -producer collective -- is not a performing act: it has no voice, -so its gender is inapplicable, and it takes no part in the -derivation. See `PerformerGenderDeriver.performer_gender`. - -Once the performer genders are derived, an optional gender -correction table CSV overrides the performer gender of the songs -it names, matched by exact title and exact artist credit, with -its performer gender column stored verbatim; see -`GenderCorrectionImporter`. - -On a successful build, two review CSV files, ``songs.csv`` and -``artists.csv``, are (re)written under the given output directory, -mirroring the stored songs and artists without their IDs; see -`CSVExporter`. A failed build leaves any existing review CSV -files untouched, matching the store rollback. +intact. See `StoreBuilder` for the pipeline, and the individual +importer, deriver, and exporter classes for the identity, +derivation, correction, and export rules. """ import argparse import csv @@ -70,7 +29,7 @@ from collections import Counter from collections.abc import Iterable, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Any, Self +from typing import Any, ClassVar, Self import sqlalchemy as sa from sqlalchemy.orm import Session @@ -89,106 +48,59 @@ from ..models import ( ) from ..utils import format_duration -ARTIST_FIELDS: dict[str, str] = { - "qid": "wikidata_qid", - "gender": "gender", - "type": "type", - "genre": "genre", - "country": "country", -} -"""The artist CSV columns mapped to the Artist attributes.""" - class BuildError(Exception): """An error that fails the build.""" -def parse_args(argv: list[str] | None) -> argparse.Namespace: - """Parse the command-line arguments. - - :param argv: The command-line arguments, or None for - ``sys.argv``. - :return: The parsed arguments. - """ - parser: argparse.ArgumentParser = argparse.ArgumentParser( - description="Rebuild the SQLite working store from the" - " committed inputs.") - parser.add_argument( - "chart_csv", type=Path, - help="the year-end chart CSV file") - parser.add_argument( - "derived_dir", type=Path, - help="the output directory for the review CSV files") - parser.add_argument( - "--lyrics-dir", type=Path, default=None, - help="the lyrics cache directory to load") - parser.add_argument( - "--wikidata-csv", type=Path, default=None, - help="the Wikidata artist snapshot CSV file to apply") - parser.add_argument( - "--codings", type=Path, default=None, - help="the settled coding table CSV file to import") - parser.add_argument( - "--groups", type=Path, default=None, - help="the settled code group table CSV file to import") - parser.add_argument( - "--gender-corrections", type=Path, default=None, - help="the gender correction table CSV file to apply") - parser.add_argument( - "--patterns", type=Path, default=None, - help="the pattern definition table CSV file to import") - parser.add_argument( - "--annotations", type=Path, default=None, - help="the settled pattern annotation table CSV file to" - " import") - return parser.parse_args(argv) - - class SongImporter: - """The song-import job: loads the chart CSV into songs and - chart entries.""" + """The song-import job: assigns every distinct chart song its + deterministic ID and records its chart appearances. - YEARS: Sequence[int] = range(2016, 2026) + A song repeated across the rows is stored once, matched by its + identity key (see `song_identity`); every row yields one chart + entry. The stored title is the raw title; the stored artist + credit is the canonical credit from the identity key. The + songs take the IDs 1, 2, 3, ... in the first-occurrence row + order. + """ + + __YEARS: ClassVar[Sequence[int]] = range(2016, 2026) """The expected chart years.""" - RANKS_PER_YEAR: int = 100 + __RANKS_PER_YEAR: ClassVar[int] = 100 """The expected number of ranks on the chart of each year.""" - CANONICAL_ARTIST_CREDITS: dict[str, str] = { + __CANONICAL_ARTIST_CREDITS: ClassVar[dict[str, str]] = { "benny blanco, Halsey & Khalid": "Benny Blanco, Halsey" " & Khalid", } """The canonical artist credit spellings, keyed by a variant credit string.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, chart_csv: Path) -> None: """Initialize the importer. :param session: The database session. + :param chart_csv: The chart CSV file with the columns + year, rank, title, and artist. """ self.__session: Session = session + self.__chart_csv: Path = chart_csv self.__songs: dict[tuple[str, str], Song] = {} - def import_songs(self, path: Path) -> None: + def run(self) -> None: """Load the chart CSV into songs and chart entries. - A song repeated across the rows is stored once, matched by - its identity key (see `song_identity`); every row yields - one chart entry. The stored title is the raw title; the - stored artist credit is the canonical credit from the - identity key. The songs take the IDs 1, 2, 3, ... in the - first-occurrence row order. When the method returns, the - imported songs and chart entries are queryable in the - session. + When the method returns, the imported songs and chart + entries are queryable in the session. - :param path: The chart CSV file with the columns year, - rank, title, and artist. :return: None. :raises BuildError: When the chart entries do not cover - each of ``YEARS`` and each rank from 1 to - ``RANKS_PER_YEAR`` exactly once. + each expected year and rank exactly once. :raises OSError: When the file cannot be read. """ counts: Counter[tuple[int, int]] = Counter() - with open(path, encoding="utf-8", newline="") as file: + with open(self.__chart_csv, encoding="utf-8", + newline="") as file: row: dict[str, str] for row in csv.DictReader(file): key: tuple[str, str] = self.song_identity( @@ -225,8 +137,8 @@ class SongImporter: present, or a pair is duplicated. """ expected: set[tuple[int, int]] = { - (year, rank) for year in cls.YEARS - for rank in range(1, cls.RANKS_PER_YEAR + 1)} + (year, rank) for year in cls.__YEARS + for rank in range(1, cls.__RANKS_PER_YEAR + 1)} actual: set[tuple[int, int]] = set(counts) violations: list[str] = [] year: int @@ -250,7 +162,7 @@ class SongImporter: """Compute the identity key of a chart row. The key pairs the raw title with the artist credit, - canonicalized through ``CANONICAL_ARTIST_CREDITS``; a + canonicalized through the canonical artist credit table; a credit absent from the table maps to itself. Two chart rows denote the same song iff their identity keys are equal. @@ -260,32 +172,43 @@ class SongImporter: :return: The identity key: the raw title paired with the canonical artist credit. """ - return title, SongImporter.CANONICAL_ARTIST_CREDITS.get( + return title, SongImporter.__CANONICAL_ARTIST_CREDITS.get( credit, credit) class ArtistImporter: - """The artist-import job: parses the stored songs' artist - credits into artists and song-artist credits.""" + """The artist-import job: resolves every credited artist name + to a deduplicated artist row and records the song-artist + credits. - FEATURING_PATTERN: re.Pattern[str] = re.compile( + An artist parsed out of a credit (see `parse_artist_credit`) is + matched against the known artists by its identity key (see + `resolve_artist_identity`); a newly seen one takes the ID + following the known artists, assigned in first-seen order + across the songs, and its stored name is the resolved stored + spelling. An artist duplicated within one song's credit, by + its identity key, is kept only at its first occurrence within + that credit, with a warning to the standard error. + """ + + __FEATURING_PATTERN: ClassVar[re.Pattern[str]] = re.compile( r" featuring | feat\. ", re.IGNORECASE) """The pattern splitting the primary and featured sides.""" - DELIMITER_PATTERN: re.Pattern[str] = re.compile( + __DELIMITER_PATTERN: ClassVar[re.Pattern[str]] = re.compile( r", | & | \+ | / |(?i: and | x | with )") """The pattern splitting the artist names within a side.""" - COLON_PATTERN: re.Pattern[str] = re.compile(r": ") + __COLON_PATTERN: ClassVar[re.Pattern[str]] = re.compile(r": ") """The pattern separating a group prefix from its members in a ": " credit.""" - PAREN_MEMBERS_PATTERN: re.Pattern[str] = re.compile( + __PAREN_MEMBERS_PATTERN: ClassVar[re.Pattern[str]] = re.compile( r"^[^(]+ \((?P[^()]+)\)$") """The pattern separating a group name from its members in a " ()" credit spanning the whole credit.""" - DUET_WITH_PATTERN: re.Pattern[str] = re.compile( + __DUET_WITH_PATTERN: ClassVar[re.Pattern[str]] = re.compile( r" Duet With ", re.IGNORECASE) """The pattern normalizing the "Duet With" co-billing connector to the plain "with" delimiter.""" - PROTECTED_ARTIST_NAMES: tuple[str, ...] = ( + __PROTECTED_ARTIST_NAMES: ClassVar[tuple[str, ...]] = ( "Tyler, The Creator", "Lil Nas X", "Tones And I", @@ -293,7 +216,8 @@ class ArtistImporter: """The exact artist names guarded from the delimiter splitting, because each contains a delimiter word or punctuation as part of the name itself.""" - EXCEPTION_CREDITS: dict[str, list[tuple[str, Role]]] = { + __EXCEPTION_CREDITS: ClassVar[dict[str, list[tuple[str, Role]]]] \ + = { "SpotemGottem Featuring Pooh Shiesty Or DaBaby": [ ("SpotemGottem", Role.PRIMARY), ("Pooh Shiesty", Role.FEATURED), @@ -312,7 +236,7 @@ class ArtistImporter: """The single-credit exceptions parsed by an explicit lookup rather than by the general rules, because the credit text alone does not spell out the correct member split.""" - CANONICAL_ARTIST_NAMES: dict[str, str] = { + __CANONICAL_ARTIST_NAMES: ClassVar[dict[str, str]] = { "beyonce": "Beyoncé", "5 seconds of summer": "5 Seconds of Summer", "a boogie wit da hoodie": "A Boogie wit da Hoodie", @@ -355,26 +279,20 @@ class ArtistImporter: self.__session: Session = session self.__artists: dict[str, Artist] = {} - def import_artists(self) -> None: + def run(self) -> None: """Parse the stored songs' credits into artists and song-artist credits. Reads the songs back from the database in ``Song.id`` order, including any songs pending in the same session, and for - each song parses ``Song.artist_credit`` (see - `parse_artist_credit`). An artist parsed out of a credit is - matched against the known artists by its identity key (see - `resolve_artist_identity`); a newly seen one takes the ID - following the known artists, keyed by its identity key, - assigned in first-seen order across the songs, and its - stored name is the resolved stored spelling. An artist - duplicated within one song's credit, by its identity key, - is kept only at its first occurrence within that credit, - with a warning to the standard error. When the method + each song parses ``Song.artist_credit``. When the method returns, the imported artists and credits are queryable in the session. :return: None. + :raises BuildError: When a parsed credit has no primary + artist or contains a blank artist name (see + `__check_parsed_credit`). """ song: Song for song in self.__session.scalars( @@ -444,7 +362,7 @@ class ArtistImporter: def parse_artist_credit(credit: str) -> list[tuple[str, Role]]: """Parse a combined artist credit into artists and roles. - A credit listed in ``EXCEPTION_CREDITS`` is looked up + A credit listed as a single-credit exception is looked up verbatim, because its correct split is not derivable from the credit text alone. Otherwise the credit first reduces to an effective credit: a ": " prefix @@ -459,7 +377,7 @@ class ArtistImporter: names on the delimiters ", ", " & ", " + ", " / " (literally) and " and ", " x ", " with " (case-insensitively), except for the names listed in - ``PROTECTED_ARTIST_NAMES``, which are never split even + ``__PROTECTED_ARTIST_NAMES``, which are never split even though each contains a delimiter word or punctuation. Known limitation: a compound act name that contains one of @@ -471,38 +389,38 @@ class ArtistImporter: side first, with the role ``Role.PRIMARY`` or ``Role.FEATURED``. """ - if credit in ArtistImporter.EXCEPTION_CREDITS: - return list(ArtistImporter.EXCEPTION_CREDITS[credit]) + if credit in ArtistImporter.__EXCEPTION_CREDITS: + return list(ArtistImporter.__EXCEPTION_CREDITS[credit]) effective: str = credit colon_match: re.Match[str] | None = \ - ArtistImporter.COLON_PATTERN.search(effective) + ArtistImporter.__COLON_PATTERN.search(effective) if colon_match is not None: effective = effective[colon_match.end():] else: paren_match: re.Match[str] | None = \ - ArtistImporter.PAREN_MEMBERS_PATTERN.match( + ArtistImporter.__PAREN_MEMBERS_PATTERN.match( effective) if paren_match is not None: effective = paren_match.group("members") - effective = ArtistImporter.DUET_WITH_PATTERN.sub( + effective = ArtistImporter.__DUET_WITH_PATTERN.sub( " with ", effective) placeholders: dict[str, str] = {} index: int protected: str for index, protected in enumerate( - ArtistImporter.PROTECTED_ARTIST_NAMES): + ArtistImporter.__PROTECTED_ARTIST_NAMES): if protected in effective: placeholder: str = f"{index}" placeholders[placeholder] = protected effective = effective.replace(protected, placeholder) - sides: list[str] = ArtistImporter.FEATURING_PATTERN.split( + sides: list[str] = ArtistImporter.__FEATURING_PATTERN.split( effective, maxsplit=1) pairs: list[tuple[str, Role]] = [] role: Role side: str for side, role in zip(sides, (Role.PRIMARY, Role.FEATURED)): token: str - for token in ArtistImporter.DELIMITER_PATTERN.split( + for token in ArtistImporter.__DELIMITER_PATTERN.split( side): name: str = ArtistImporter.__restore_protected( token.strip(), placeholders) @@ -532,9 +450,9 @@ class ArtistImporter: def resolve_artist_identity(name: str) -> tuple[str, str]: """Resolve the dedup key and the stored spelling of a name. - The name's case-folded form is looked up in - ``CANONICAL_ARTIST_NAMES`` first; when it is listed there, - the dedup key is the canonical spelling case-folded and the + The name's case-folded form is looked up in the canonical + artist name table first; when it is listed there, the + dedup key is the canonical spelling case-folded and the stored spelling is the canonical spelling, so every variant of the name, canonical or not, resolves to the same identity. Otherwise the dedup key is the name case-folded @@ -545,7 +463,7 @@ class ArtistImporter: """ folded: str = name.casefold() canonical: str | None = \ - ArtistImporter.CANONICAL_ARTIST_NAMES.get(folded) + ArtistImporter.__CANONICAL_ARTIST_NAMES.get(folded) if canonical is not None: return canonical.casefold(), canonical return folded, name @@ -553,42 +471,55 @@ class ArtistImporter: class CaptureImporter: """The capture-import job: applies the optional capture-layer - inputs onto the stored songs and artists.""" + inputs -- the lyrics cache and the Wikidata artist snapshot -- + onto the stored songs and artists. - def __init__(self, session: Session) -> None: + A None input leaves its capture layer unloaded. A given + artist snapshot applies only its non-empty cells, field by + field; the note column is ignored. + """ + + __ARTIST_FIELDS: ClassVar[dict[str, str]] = { + "qid": "wikidata_qid", + "gender": "gender", + "type": "type", + "genre": "genre", + "country": "country", + } + """The artist CSV columns mapped to the Artist attributes.""" + + def __init__(self, session: Session, lyrics_dir: Path | None, + wikidata_csv: Path | None) -> None: """Initialize the importer. :param session: The database session. - """ - self.__session: Session = session - - def import_captures(self, lyrics_dir: Path | None, - wikidata_csv: Path | None) -> None: - """Apply the optional capture-layer inputs onto the store. - - A None input leaves its capture layer unloaded. When - ``lyrics_dir`` is given, its cached lyrics files load into - the matching songs (see `__load_lyrics`). When - ``wikidata_csv`` is given, it applies onto the artist rows - (see `__apply_artist_csv`). When the method returns, the - applied changes are queryable in the session. - :param lyrics_dir: The lyrics cache directory to load, or None to skip the lyrics capture layer. :param wikidata_csv: The Wikidata artist snapshot CSV file to apply, or None to skip the artist capture layer. + """ + self.__session: Session = session + self.__lyrics_dir: Path | None = lyrics_dir + self.__wikidata_csv: Path | None = wikidata_csv + + def run(self) -> None: + """Apply the optional capture-layer inputs onto the store. + + When the method returns, the applied changes are queryable + in the session. + :return: None. :raises BuildError: When ``lyrics_dir`` does not exist, or a name in ``wikidata_csv`` matches no artist. :raises OSError: When a capture file cannot be read. """ - if lyrics_dir is not None: - if not lyrics_dir.is_dir(): + if self.__lyrics_dir is not None: + if not self.__lyrics_dir.is_dir(): raise BuildError( - f"{lyrics_dir}: no such directory") - self.__load_lyrics(lyrics_dir) - if wikidata_csv is not None: - self.__apply_artist_csv(wikidata_csv) + f"{self.__lyrics_dir}: no such directory") + self.__load_lyrics(self.__lyrics_dir) + if self.__wikidata_csv is not None: + self.__apply_artist_csv(self.__wikidata_csv) self.__session.flush() def __load_lyrics(self, directory: Path) -> None: @@ -615,8 +546,7 @@ class CaptureImporter: def __apply_artist_csv(self, path: Path) -> None: """Apply an artist attribute CSV onto the artist rows. - Artists match by exact name. Only the non-empty cells are - applied, field by field. The note column is ignored. + Artists match by exact name. :param path: The CSV file with the columns name, qid, gender, type, genre, country, and note. @@ -636,7 +566,8 @@ class CaptureImporter: f" \"{row['name']}\"") column: str attribute: str - for column, attribute in ARTIST_FIELDS.items(): + for column, attribute in \ + self.__ARTIST_FIELDS.items(): if row.get(column): setattr(artist, attribute, row[column]) @@ -647,7 +578,7 @@ class PerformerGenderDeriver: performing acts, a credited artist without an artist type taking no part.""" - MIXED: str = "mixed" + __MIXED: ClassVar[str] = "mixed" """The performer gender of a song whose performing credited artists do not all share one gender.""" @@ -658,7 +589,7 @@ class PerformerGenderDeriver: """ self.__session: Session = session - def derive_performer_genders(self) -> None: + def run(self) -> None: """Set the performer gender of every stored song. Reads the songs back from the database, including any songs @@ -685,12 +616,12 @@ class PerformerGenderDeriver: """Combine the performing artists' genders into one value. A gender that is None or empty counts as unknown. Two or - more distinct known genders give ``MIXED``, an unknown one - notwithstanding, as an unknown cannot undo a disagreement. - A single known gender shared by every given artist gives - that gender. Anything else -- a single known gender - alongside an unknown one, no known gender at all, or no - gender given at all -- gives None. + more distinct known genders give the mixed gender, an + unknown one notwithstanding, as an unknown cannot undo a + disagreement. A single known gender shared by every given + artist gives that gender. Anything else -- a single known + gender alongside an unknown one, no known gender at all, or + no gender given at all -- gives None. :param genders: The genders of the performing artists credited on one song, in any order. @@ -700,70 +631,138 @@ class PerformerGenderDeriver: values: list[str | None] = list(genders) known: set[str] = {x for x in values if x} if len(known) > 1: - return cls.MIXED + return cls.__MIXED if len(known) == 1 and all(x for x in values): return known.pop() return None -class GenderCorrectionImporter: - """The gender-correction job: overrides the derived performer - gender of the stored songs it names.""" +class ColumnCheckedImporter: + """The shared skeleton of an optional, header-checked CSV + importer: a None path leaves the layer unloaded; otherwise the + header is validated and every row imports within one flush. + A subclass supplies its required columns to the constructor and + overrides `prepare` and `import_row` for its own row-handling. + """ - COLUMNS: tuple[str, ...] = ( + def __init__(self, session: Session, path: Path | None, + columns: Sequence[str]) -> None: + """Initialize the shared importer skeleton. + + :param session: The database session. + :param path: The CSV file to import, or None to skip. + :param columns: The required column names of the file. + """ + self.__session: Session = session + self.__path: Path | None = path + self.__columns: Sequence[str] = columns + + @property + def session(self) -> Session: + """The database session.""" + return self.__session + + @property + def path(self) -> Path | None: + """The CSV file being imported, or None when skipped.""" + return self.__path + + def run(self) -> None: + """Validate the header and import every row of the file. + + A None path leaves the layer unloaded. Otherwise calls + `prepare` once, then `import_row` for every data row, and + flushes the session once every row is stored. + + :return: None. + :raises BuildError: When the file lacks a required column, + or a subclass raises importing a row. + :raises OSError: When the file cannot be read. + """ + if self.__path is None: + return + self.prepare() + with open(self.__path, encoding="utf-8", + newline="") as file: + reader: csv.DictReader[str] = csv.DictReader(file) + self.__check_columns(reader.fieldnames) + row: dict[str, str] + for row in reader: + self.import_row(row) + self.__session.flush() + + def __check_columns( + self, fieldnames: Sequence[str] | None) -> None: + """Verify the CSV file has the required columns. + + :param fieldnames: The header row of the file, or None + when the file is empty. + :return: None. + :raises BuildError: When a required column is absent. + """ + header: Sequence[str] = fieldnames or () + missing: list[str] = [ + x for x in self.__columns if x not in header] + if len(missing) > 0: + raise BuildError( + f"{self.__path}: missing column(s):" + f" {', '.join(missing)}") + + def prepare(self) -> None: + """Set up any state needed before the rows import. + + The default is a no-op; a subclass overrides this to load + database lookups needed by `import_row`. + + :return: None. + """ + return None + + def import_row(self, row: dict[str, str]) -> None: + """Import one data row. + + A subclass overrides this to store the row. + + :param row: The CSV row. + :return: None. + """ + raise NotImplementedError + + +class GenderCorrectionImporter(ColumnCheckedImporter): + """The gender-correction job: overrides the derived performer + gender of the stored songs it names, matched by exact title and + exact artist credit, with the performer gender column stored + verbatim. Run this after `PerformerGenderDeriver`, so a + correction overrides the derived value. + """ + + __COLUMNS: ClassVar[tuple[str, ...]] = ( "Title", "Artist Credit", "Performer Gender", "Note") """The required columns of the gender correction CSV file.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, path: Path | None) -> None: """Initialize the importer. :param session: The database session. + :param path: The gender correction table CSV file to + apply, or None to skip the corrections. """ - self.__session: Session = session + super().__init__(session, path, self.__COLUMNS) + self.__songs: dict[tuple[str, str], Song] = {} - def import_gender_corrections(self, path: Path | None) -> None: - """Apply the gender correction table onto the stored songs. + def prepare(self) -> None: + """Load the stored songs, keyed by title and artist credit. - A None input leaves the derived performer genders - untouched. Otherwise every row of the CSV file matches one - stored song by its exact title and exact artist credit and - sets ``Song.performer_gender`` to the row's performer - gender column verbatim; the note column is ignored. Apply - this after `PerformerGenderDeriver.derive_performer_genders` - has run, so a correction overrides the derived value. When - the method returns, the applied corrections are queryable - in the session. - - :param path: The gender correction table CSV file to apply, - or None to skip the corrections. :return: None. - :raises BuildError: When the file lacks a required column, - or a row names a song that the store does not have. - :raises OSError: When the file cannot be read. """ - if path is None: - return - songs: dict[tuple[str, str], Song] = { + self.__songs = { (x.title, x.artist_credit): x - for x in self.__session.scalars(sa.select(Song))} - with open(path, encoding="utf-8", newline="") as file: - reader: csv.DictReader[str] = csv.DictReader(file) - self.__check_columns(path, reader.fieldnames) - row: dict[str, str] - for row in reader: - self.__apply_correction(path, songs, row) - self.__session.flush() + for x in self.session.scalars(sa.select(Song))} - @staticmethod - def __apply_correction(path: Path, - songs: dict[tuple[str, str], Song], - row: dict[str, str]) -> None: + def import_row(self, row: dict[str, str]) -> None: """Apply one gender correction row. - :param path: The gender correction CSV file, for the error - message. - :param songs: The stored songs, keyed by the title and the - artist credit. :param row: The gender correction CSV row. :return: None. :raises BuildError: When the row names a song that the @@ -771,93 +770,49 @@ class GenderCorrectionImporter: """ key: tuple[str, str] = ( row["Title"], row["Artist Credit"]) - song: Song | None = songs.get(key) + song: Song | None = self.__songs.get(key) if song is None: raise BuildError( - f"{path}: no song \"{row['Title']}\" by" + f"{self.path}: no song \"{row['Title']}\" by" f" \"{row['Artist Credit']}\"") song.performer_gender = row["Performer Gender"] - @classmethod - def __check_columns(cls, path: Path, - fieldnames: Sequence[str] | None) -> None: - """Verify the gender correction CSV file has the required - columns. - :param path: The gender correction CSV file. - :param fieldnames: The header row of the file, or None - when the file is empty. - :return: None. - :raises BuildError: When a required column is absent. - """ - header: Sequence[str] = fieldnames or () - missing: list[str] = [ - x for x in cls.COLUMNS if x not in header] - if len(missing) > 0: - raise BuildError( - f"{path}: missing column(s): {', '.join(missing)}") - - -class CodingImporter: +class CodingImporter(ColumnCheckedImporter): """The coding-import job: loads the settled coding table onto - the stored songs.""" + the stored songs, matched by exact title and exact artist + credit, with the quote column stored verbatim; a quote carries + the lyric line-break convention ``" / "`` where the lyric has a + line break, and an empty quote column stores an empty string. + """ - COLUMNS: tuple[str, ...] = ( + __COLUMNS: ClassVar[tuple[str, ...]] = ( "Song", "Artist Credit", "Keyword", "Quote") """The required columns of the coding CSV file.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, path: Path | None) -> None: """Initialize the importer. :param session: The database session. - """ - self.__session: Session = session - - def import_codings(self, path: Path | None) -> None: - """Load the settled coding table onto the stored songs. - - A None input leaves the coding unloaded. Otherwise every - row of the CSV file yields one coding of the song named by - its title and artist credit, with the quote column stored - verbatim; a quote carries the lyric line-break convention - ``" / "`` where the lyric has a line break, and an empty - quote column stores an empty string. When the method - returns, the imported codings are queryable in the - session. - :param path: The settled coding table CSV file to import, or None to skip the coding. - :return: None. - :raises BuildError: When the file lacks a required column, - a row names a song that the store does not have, or two - rows name the same song and keyword. - :raises OSError: When the file cannot be read. """ - if path is None: - return - songs: dict[tuple[str, str], Song] = { - (x.title, x.artist_credit): x - for x in self.__session.scalars(sa.select(Song))} - seen: set[tuple[int, str]] = set() - with open(path, encoding="utf-8", newline="") as file: - reader: csv.DictReader[str] = csv.DictReader(file) - self.__check_columns(path, reader.fieldnames) - row: dict[str, str] - for row in reader: - self.__import_coding(path, songs, seen, row) - self.__session.flush() + super().__init__(session, path, self.__COLUMNS) + self.__songs: dict[tuple[str, str], Song] = {} + self.__seen: set[tuple[int, str]] = set() - def __import_coding(self, path: Path, - songs: dict[tuple[str, str], Song], - seen: set[tuple[int, str]], - row: dict[str, str]) -> None: + def prepare(self) -> None: + """Load the stored songs, keyed by title and artist credit. + + :return: None. + """ + self.__songs = { + (x.title, x.artist_credit): x + for x in self.session.scalars(sa.select(Song))} + + def import_row(self, row: dict[str, str]) -> None: """Store one coding row. - :param path: The coding CSV file, for the error messages. - :param songs: The stored songs, keyed by the title and the - artist credit. - :param seen: The (song ID, keyword) pairs already stored, - updated with the pair of this row. :param row: The coding CSV row. :return: None. :raises BuildError: When the row names a song that the @@ -865,91 +820,45 @@ class CodingImporter: earlier row. """ key: tuple[str, str] = (row["Song"], row["Artist Credit"]) - song: Song | None = songs.get(key) + song: Song | None = self.__songs.get(key) if song is None: raise BuildError( - f"{path}: no song \"{row['Song']}\" by" + f"{self.path}: no song \"{row['Song']}\" by" f" \"{row['Artist Credit']}\"") coding_key: tuple[int, str] = (song.id, row["Keyword"]) - if coding_key in seen: + if coding_key in self.__seen: raise BuildError( - f"{path}: duplicated coding: \"{row['Song']}\" by" - f" \"{row['Artist Credit']}\", keyword" + f"{self.path}: duplicated coding: \"{row['Song']}\"" + f" by \"{row['Artist Credit']}\", keyword" f" \"{row['Keyword']}\"") - seen.add(coding_key) - self.__session.add(Coding( + self.__seen.add(coding_key) + self.session.add(Coding( song=song, keyword=row["Keyword"], quotes=row["Quote"])) - @classmethod - def __check_columns(cls, path: Path, - fieldnames: Sequence[str] | None) -> None: - """Verify the coding CSV file has the required columns. - :param path: The coding CSV file. - :param fieldnames: The header row of the file, or None when - the file is empty. - :return: None. - :raises BuildError: When a required column is absent. - """ - header: Sequence[str] = fieldnames or () - missing: list[str] = [ - x for x in cls.COLUMNS if x not in header] - if len(missing) > 0: - raise BuildError( - f"{path}: missing column(s): {', '.join(missing)}") - - -class GroupImporter: +class GroupImporter(ColumnCheckedImporter): """The group-import job: loads the settled code group table - into the working store.""" + into the working store, with the group name, the keyword, and + the integer vote count stored verbatim.""" - COLUMNS: tuple[str, ...] = ("Group", "Keyword", "Votes") + __COLUMNS: ClassVar[tuple[str, ...]] = ( + "Group", "Keyword", "Votes") """The required columns of the group CSV file.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, path: Path | None) -> None: """Initialize the importer. :param session: The database session. - """ - self.__session: Session = session - - def import_groups(self, path: Path | None) -> None: - """Load the settled code group table into the store. - - A None input leaves the groups unloaded. Otherwise every - row of the CSV file yields one member keyword of one - group, with the group name, the keyword, and the integer - vote count stored verbatim. When the method returns, the - imported groups are queryable in the session. - :param path: The settled code group table CSV file to import, or None to skip the groups. - :return: None. - :raises BuildError: When the file lacks a required - column, a votes field is not an integer, or two rows - name the same group and keyword. - :raises OSError: When the file cannot be read. """ - if path is None: - return - seen: set[tuple[str, str]] = set() - with open(path, encoding="utf-8", newline="") as file: - reader: csv.DictReader[str] = csv.DictReader(file) - self.__check_columns(path, reader.fieldnames) - row: dict[str, str] - for row in reader: - self.__import_group_row(path, seen, row) - self.__session.flush() + super().__init__(session, path, self.__COLUMNS) + self.__seen: set[tuple[str, str]] = set() - def __import_group_row(self, path: Path, - seen: set[tuple[str, str]], - row: dict[str, str]) -> None: + def import_row(self, row: dict[str, str]) -> None: """Store one group member row. - :param path: The group CSV file, for the error messages. - :param seen: The (group, keyword) pairs already stored, - updated with the pair of this row. :param row: The group CSV row. :return: None. :raises BuildError: When the votes field is not an @@ -957,196 +866,103 @@ class GroupImporter: row. """ key: tuple[str, str] = (row["Group"], row["Keyword"]) - if key in seen: + if key in self.__seen: raise BuildError( - f"{path}: duplicated group member: group" + f"{self.path}: duplicated group member: group" f" \"{row['Group']}\", keyword" f" \"{row['Keyword']}\"") - seen.add(key) + self.__seen.add(key) votes: int try: votes = int(row["Votes"]) except ValueError as error: raise BuildError( - f"{path}: group \"{row['Group']}\", keyword" + f"{self.path}: group \"{row['Group']}\", keyword" f" \"{row['Keyword']}\": votes" f" \"{row['Votes']}\" is not an integer" ) from error - self.__session.add(CodeGroup( + self.session.add(CodeGroup( group=row["Group"], keyword=row["Keyword"], votes=votes)) - @classmethod - def __check_columns(cls, path: Path, - fieldnames: Sequence[str] | None) -> None: - """Verify the group CSV file has the required columns. - :param path: The group CSV file. - :param fieldnames: The header row of the file, or None - when the file is empty. - :return: None. - :raises BuildError: When a required column is absent. - """ - header: Sequence[str] = fieldnames or () - missing: list[str] = [ - x for x in cls.COLUMNS if x not in header] - if len(missing) > 0: - raise BuildError( - f"{path}: missing column(s): {', '.join(missing)}") - - -class PatternImporter: +class PatternImporter(ColumnCheckedImporter): """The pattern-import job: loads the pattern definition table - into the working store.""" + into the working store, with the pattern ID, the group, the + name, and the description stored verbatim.""" - COLUMNS: tuple[str, ...] = ( + __COLUMNS: ClassVar[tuple[str, ...]] = ( "Pattern", "Group", "Name", "Description") """The required columns of the pattern CSV file.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, path: Path | None) -> None: """Initialize the importer. :param session: The database session. - """ - self.__session: Session = session - - def import_patterns(self, path: Path | None) -> None: - """Load the pattern definition table into the store. - - A None input leaves the patterns unloaded. Otherwise - every row of the CSV file stores one pattern, with the - pattern ID, the group, the name, and the description - stored verbatim. When the method returns, the imported - patterns are queryable in the session. - :param path: The pattern definition table CSV file to import, or None to skip the patterns. - :return: None. - :raises BuildError: When the file lacks a required - column, or two rows name the same pattern ID. - :raises OSError: When the file cannot be read. """ - if path is None: - return - seen: set[str] = set() - with open(path, encoding="utf-8", newline="") as file: - reader: csv.DictReader[str] = csv.DictReader(file) - self.__check_columns(path, reader.fieldnames) - row: dict[str, str] - for row in reader: - self.__import_pattern_row(path, seen, row) - self.__session.flush() + super().__init__(session, path, self.__COLUMNS) + self.__seen: set[str] = set() - def __import_pattern_row(self, path: Path, seen: set[str], - row: dict[str, str]) -> None: + def import_row(self, row: dict[str, str]) -> None: """Store one pattern row. - :param path: The pattern CSV file, for the error messages. - :param seen: The pattern IDs already stored, updated with - this row's pattern ID. :param row: The pattern CSV row. :return: None. :raises BuildError: When the pattern ID repeats an earlier row. """ - if row["Pattern"] in seen: + if row["Pattern"] in self.__seen: raise BuildError( - f"{path}: duplicated pattern \"{row['Pattern']}\"") - seen.add(row["Pattern"]) - self.__session.add(Pattern( + f"{self.path}: duplicated pattern" + f" \"{row['Pattern']}\"") + self.__seen.add(row["Pattern"]) + self.session.add(Pattern( pattern=row["Pattern"], group=row["Group"], name=row["Name"], description=row["Description"])) - @classmethod - def __check_columns(cls, path: Path, - fieldnames: Sequence[str] | None) -> None: - """Verify the pattern CSV file has the required columns. - :param path: The pattern CSV file. - :param fieldnames: The header row of the file, or None - when the file is empty. - :return: None. - :raises BuildError: When a required column is absent. - """ - header: Sequence[str] = fieldnames or () - missing: list[str] = [ - x for x in cls.COLUMNS if x not in header] - if len(missing) > 0: - raise BuildError( - f"{path}: missing column(s): {', '.join(missing)}") - - -class AnnotationImporter: +class AnnotationImporter(ColumnCheckedImporter): """The annotation-import job: loads the settled pattern - annotation table onto the stored songs.""" + annotation table onto the stored songs, linking the song named + by its title and artist credit to the stored pattern named by + its pattern ID, with the integer vote count stored verbatim. + Import the pattern table first (see `PatternImporter`), so + every named pattern is already stored. + """ - COLUMNS: tuple[str, ...] = ( + __COLUMNS: ClassVar[tuple[str, ...]] = ( "Song", "Artist Credit", "Pattern", "Votes") """The required columns of the annotation CSV file.""" - def __init__(self, session: Session) -> None: + def __init__(self, session: Session, path: Path | None) -> None: """Initialize the importer. :param session: The database session. - """ - self.__session: Session = session - - def import_annotations(self, path: Path | None) -> None: - """Load the settled pattern annotation table onto the - stored songs. - - A None input leaves the annotations unloaded. Otherwise - every row of the CSV file yields one annotation linking - the song named by its title and artist credit to the - stored pattern named by its pattern ID, with the integer - vote count stored verbatim. Import the pattern table - first (see `PatternImporter`), so every named pattern is - already stored. When the method returns, the imported - annotations are queryable in the session. - :param path: The settled pattern annotation table CSV file to import, or None to skip the annotations. - :return: None. - :raises BuildError: When the file lacks a required - column, a row names a song that the store does not - have, a row names a pattern that the store does not - have, a votes field is not an integer, or two rows - name the same song and pattern. - :raises OSError: When the file cannot be read. """ - if path is None: - return - songs: dict[tuple[str, str], Song] = { - (x.title, x.artist_credit): x - for x in self.__session.scalars(sa.select(Song))} - patterns: dict[str, Pattern] = { - x.pattern: x - for x in self.__session.scalars(sa.select(Pattern))} - seen: set[tuple[int, str]] = set() - with open(path, encoding="utf-8", newline="") as file: - reader: csv.DictReader[str] = csv.DictReader(file) - self.__check_columns(path, reader.fieldnames) - row: dict[str, str] - for row in reader: - self.__import_annotation( - path, songs, patterns, seen, row) - self.__session.flush() + super().__init__(session, path, self.__COLUMNS) + self.__songs: dict[tuple[str, str], Song] = {} + self.__patterns: dict[str, Pattern] = {} + self.__seen: set[tuple[int, str]] = set() - def __import_annotation( - self, path: Path, songs: dict[tuple[str, str], Song], - patterns: dict[str, Pattern], - seen: set[tuple[int, str]], - row: dict[str, str]) -> None: + def prepare(self) -> None: + """Load the stored songs and patterns for the row lookups. + + :return: None. + """ + self.__songs = { + (x.title, x.artist_credit): x + for x in self.session.scalars(sa.select(Song))} + self.__patterns = { + x.pattern: x + for x in self.session.scalars(sa.select(Pattern))} + + def import_row(self, row: dict[str, str]) -> None: """Store one annotation row. - :param path: The annotation CSV file, for the error - messages. - :param songs: The stored songs, keyed by the title and - the artist credit. - :param patterns: The stored patterns, keyed by the - pattern ID. - :param seen: The (song ID, pattern ID) pairs already - stored, updated with the pair of this row. :param row: The annotation CSV row. :return: None. :raises BuildError: When the row names a song that the @@ -1155,55 +971,36 @@ class AnnotationImporter: its song and pattern repeat an earlier row. """ key: tuple[str, str] = (row["Song"], row["Artist Credit"]) - song: Song | None = songs.get(key) + song: Song | None = self.__songs.get(key) if song is None: raise BuildError( - f"{path}: no song \"{row['Song']}\" by" + f"{self.path}: no song \"{row['Song']}\" by" f" \"{row['Artist Credit']}\"") - pattern: Pattern | None = patterns.get(row["Pattern"]) + pattern: Pattern | None = self.__patterns.get(row["Pattern"]) if pattern is None: raise BuildError( - f"{path}: no pattern \"{row['Pattern']}\"") + f"{self.path}: no pattern \"{row['Pattern']}\"") annotation_key: tuple[int, str] = ( song.id, pattern.pattern) - if annotation_key in seen: + if annotation_key in self.__seen: raise BuildError( - f"{path}: duplicated annotation: \"{row['Song']}\"" - f" by \"{row['Artist Credit']}\", pattern" - f" \"{row['Pattern']}\"") - seen.add(annotation_key) + f"{self.path}: duplicated annotation:" + f" \"{row['Song']}\" by \"{row['Artist Credit']}\"," + f" pattern \"{row['Pattern']}\"") + self.__seen.add(annotation_key) votes: int try: votes = int(row["Votes"]) except ValueError as error: raise BuildError( - f"{path}: \"{row['Song']}\" by" + f"{self.path}: \"{row['Song']}\" by" f" \"{row['Artist Credit']}\", pattern" f" \"{row['Pattern']}\": votes" f" \"{row['Votes']}\" is not an integer" ) from error - self.__session.add(Annotation( + self.session.add(Annotation( song=song, pattern=pattern, votes=votes)) - @classmethod - def __check_columns(cls, path: Path, - fieldnames: Sequence[str] | None) -> None: - """Verify the annotation CSV file has the required - columns. - - :param path: The annotation CSV file. - :param fieldnames: The header row of the file, or None - when the file is empty. - :return: None. - :raises BuildError: When a required column is absent. - """ - header: Sequence[str] = fieldnames or () - missing: list[str] = [ - x for x in cls.COLUMNS if x not in header] - if len(missing) > 0: - raise BuildError( - f"{path}: missing column(s): {', '.join(missing)}") - @dataclass class StoreCounts: @@ -1213,14 +1010,6 @@ class StoreCounts: """The number of the songs.""" artists: int """The number of the artists.""" - codings: int - """The number of the settled codings.""" - groups: int - """The number of the settled code group members.""" - patterns: int - """The number of the stored patterns.""" - annotations: int - """The number of the settled pattern annotations.""" @classmethod def get_instance(cls, session: Session) -> Self: @@ -1239,37 +1028,17 @@ class StoreCounts: songs=count( sa.select(sa.func.count()).select_from(Song)), artists=count( - sa.select(sa.func.count()).select_from(Artist)), - codings=count( - sa.select(sa.func.count()).select_from(Coding)), - groups=count( - sa.select(sa.func.count()).select_from(CodeGroup)), - patterns=count( - sa.select(sa.func.count()).select_from(Pattern)), - annotations=count( - sa.select(sa.func.count()) - .select_from(Annotation))) - - -def reset_store(session: Session) -> None: - """Delete all the rows from every table of the store. - - :param session: The database session. - :return: None. - """ - model: type[Base] - for model in (CodeGroup, Coding, Annotation, SongArtist, - ChartEntry, Song, Pattern, Artist): - session.execute(sa.delete(model)) + sa.select(sa.func.count()).select_from(Artist))) class CSVExporter: - """Writes the review CSV files mirroring the working store.""" + """Writes the review CSV files mirroring the working store, + without a song or an artist ID, on a successful build.""" - __SONGS_HEADER: tuple[str, ...] = ( + __SONGS_HEADER: ClassVar[tuple[str, ...]] = ( "Title", "Artists", "Positions", "Performer Gender") """The header row of ``songs.csv``, for human readers.""" - __ARTISTS_HEADER: tuple[str, ...] = ( + __ARTISTS_HEADER: ClassVar[tuple[str, ...]] = ( "Name", "Wikidata QID", "Gender", "Type", "Genre", "Country", "Songs") """The header row of ``artists.csv``, for human readers.""" @@ -1402,6 +1171,139 @@ class CSVExporter: return rows +class StoreBuilder: + """The orchestrator of a full working store rebuild: resets the + store, runs the importers and derivers in the pipeline order, + and writes the review CSV files, all in one transaction + committed only on success.""" + + def __init__( + self, chart_csv: Path, derived_dir: Path, + lyrics_dir: Path | None, wikidata_csv: Path | None, + codings: Path | None, groups: Path | None, + gender_corrections: Path | None, + patterns: Path | None, + annotations: Path | None) -> None: + """Set up the builder of the working store rebuild. + + :param chart_csv: The year-end chart CSV file. + :param derived_dir: The output directory for the review + CSV files. + :param lyrics_dir: The lyrics cache directory to load, or + None to skip it. + :param wikidata_csv: The Wikidata artist snapshot CSV file + to apply, or None to skip it. + :param codings: The settled coding table CSV file to + import, or None to skip it. + :param groups: The settled code group table CSV file to + import, or None to skip it. + :param gender_corrections: The gender correction table CSV + file to apply, or None to skip it. + :param patterns: The pattern definition table CSV file to + import, or None to skip it. + :param annotations: The settled pattern annotation table + CSV file to import, or None to skip it. + """ + self.__chart_csv: Path = chart_csv + self.__derived_dir: Path = derived_dir + self.__lyrics_dir: Path | None = lyrics_dir + self.__wikidata_csv: Path | None = wikidata_csv + self.__codings: Path | None = codings + self.__groups: Path | None = groups + self.__gender_corrections: Path | None = gender_corrections + self.__patterns: Path | None = patterns + self.__annotations: Path | None = annotations + + def run(self) -> StoreCounts: + """Rebuild the working store from the configured inputs. + + :return: The row counts of the rebuilt working store. + :raises BuildError: When an input is malformed, as + detailed on the importer and deriver classes. + :raises OSError: When an input or output file cannot be + read or written. + """ + Base.metadata.create_all(ds.engine) + session: Session = ds.get_db() + counts: StoreCounts + try: + self.__reset_store(session) + SongImporter(session, self.__chart_csv).run() + ArtistImporter(session).run() + CaptureImporter( + session, self.__lyrics_dir, + self.__wikidata_csv).run() + PerformerGenderDeriver(session).run() + GenderCorrectionImporter( + session, self.__gender_corrections).run() + CodingImporter(session, self.__codings).run() + GroupImporter(session, self.__groups).run() + PatternImporter(session, self.__patterns).run() + AnnotationImporter(session, self.__annotations).run() + counts = StoreCounts.get_instance(session) + CSVExporter(session, self.__derived_dir).write() + session.commit() + except (OSError, BuildError): + session.rollback() + raise + finally: + session.close() + return counts + + @staticmethod + def __reset_store(session: Session) -> None: + """Delete all the rows from every table of the store. + + :param session: The database session. + :return: None. + """ + model: type[Base] + for model in (CodeGroup, Coding, Annotation, SongArtist, + ChartEntry, Song, Pattern, Artist): + session.execute(sa.delete(model)) + + +def parse_args(argv: list[str] | None) -> argparse.Namespace: + """Parse the command-line arguments. + + :param argv: The command-line arguments, or None for + ``sys.argv``. + :return: The parsed arguments. + """ + parser: argparse.ArgumentParser = argparse.ArgumentParser( + description="Rebuild the SQLite working store from the" + " committed inputs.") + parser.add_argument( + "chart_csv", type=Path, + help="the year-end chart CSV file") + parser.add_argument( + "derived_dir", type=Path, + help="the output directory for the review CSV files") + parser.add_argument( + "--lyrics-dir", type=Path, default=None, + help="the lyrics cache directory to load") + parser.add_argument( + "--wikidata-csv", type=Path, default=None, + help="the Wikidata artist snapshot CSV file to apply") + parser.add_argument( + "--codings", type=Path, default=None, + help="the settled coding table CSV file to import") + parser.add_argument( + "--groups", type=Path, default=None, + help="the settled code group table CSV file to import") + parser.add_argument( + "--gender-corrections", type=Path, default=None, + help="the gender correction table CSV file to apply") + parser.add_argument( + "--patterns", type=Path, default=None, + help="the pattern definition table CSV file to import") + parser.add_argument( + "--annotations", type=Path, default=None, + help="the settled pattern annotation table CSV file to" + " import") + return parser.parse_args(argv) + + def main(argv: list[str] | None = None) -> int: """Rebuild the SQLite working store from the inputs. @@ -1411,37 +1313,17 @@ def main(argv: list[str] | None = None) -> int: """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) - Base.metadata.create_all(ds.engine) - session: Session = ds.get_db() - counts: StoreCounts try: - reset_store(session) - SongImporter(session).import_songs(args.chart_csv) - ArtistImporter(session).import_artists() - CaptureImporter(session).import_captures( - args.lyrics_dir, args.wikidata_csv) - PerformerGenderDeriver(session).derive_performer_genders() - GenderCorrectionImporter(session).import_gender_corrections( - args.gender_corrections) - CodingImporter(session).import_codings(args.codings) - GroupImporter(session).import_groups(args.groups) - PatternImporter(session).import_patterns(args.patterns) - AnnotationImporter(session).import_annotations( - args.annotations) - counts = StoreCounts.get_instance(session) - CSVExporter(session, args.derived_dir).write() - session.commit() + counts: StoreCounts = StoreBuilder( + args.chart_csv, args.derived_dir, args.lyrics_dir, + args.wikidata_csv, args.codings, args.groups, + args.gender_corrections, args.patterns, + args.annotations).run() except (OSError, BuildError) as error: - session.rollback() print(f"error: {error}", file=sys.stderr) return 1 - finally: - session.close() elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. {counts.songs} songs/{counts.artists} artists" - f"/{counts.codings} codings/{counts.groups} group" - f" members/{counts.patterns} patterns" - f"/{counts.annotations} annotations. {elapsed}" - " elapsed.", + print(f"Done. {counts.songs} songs/{counts.artists} artists." + f" {elapsed} elapsed.", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/cluster_keywords.py b/tools/src/pop_fem_audit_tools/commands/cluster_keywords.py index 5a18d48..59a38d2 100644 --- a/tools/src/pop_fem_audit_tools/commands/cluster_keywords.py +++ b/tools/src/pop_fem_audit_tools/commands/cluster_keywords.py @@ -24,8 +24,7 @@ keyword set for ``export-llm-input --extras`` is written as a JSON file holding the group name keywords plus every extra a-priori keyword the caller gives with the repeatable ``--extra-keyword`` command-line option, as -:attr:`KeywordsToMerge.KEYWORDS_TO_MERGE_JSON`; with no -``--extra-keyword``, it holds the group names alone. No default +:attr:`KeywordsToMerge.KEYWORDS_TO_MERGE_JSON`. No default extra keyword is ever injected; the caller supplies each one consciously. Finally, the command-line choices and the environment that produced the numbers -- neither recoverable from @@ -160,15 +159,8 @@ class KeywordPooler: if line.strip() == "": continue record: Any = json.loads(line) - if not isinstance(record, dict) or "id" not in record: - raise ValueError( - f"{path}: record without \"id\": {line}") if "error" in record: continue - if "text" not in record: - raise ValueError( - f"{path}: id {record['id']}: record without" - " \"text\" or \"error\"") song_id: int = cls.__parse_song_id(record["id"], path) try: keywords: Any = json.loads( @@ -297,7 +289,8 @@ class KeywordGroups: class KeywordClusterer: """The clusterer of the pooled keywords into coding groups.""" - DEFAULT_MODEL: str = "sentence-transformers/all-mpnet-base-v2" + DEFAULT_MODEL: ClassVar[str] \ + = "sentence-transformers/all-mpnet-base-v2" """The sentence embedding model used when the caller names none.""" @@ -668,11 +661,7 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: parser.add_argument( "output_dir", type=Path, help="the output directory, created if missing, that" - f" receives {PooledKeywords.SOURCE_KEYWORDS_TXT}," - f" {KeywordGroups.RESULT_KEYWORDS_TXT}," - f" {KeywordGroups.RESULT_GROUPS_CSV}," - f" {KeywordsToMerge.KEYWORDS_TO_MERGE_JSON}," - f" and {RunMeta.META_JSON}") + " receives the run's output artifacts") parser.add_argument( "--model", default=model, help=f"the sentence embedding model (default \"{model}\")") @@ -698,16 +687,9 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> int: """Pool the two tagging runs' keywords and cluster them. - Writes the five fixed-named artifacts under the output - directory, creating it (with parents) if it does not exist: - the pooled keyword text file; then the group membership CSV - file, holding the clustering result alone; the group name - keyword text file, holding the same group names as a readable - list; the coding keyword set JSON file, holding the group - names plus every extra keyword given via ``--extra-keyword``; - and the run metadata JSON file, recording the command-line - choices and the environment. Each file is written as soon as - its content is computed, so when the input is rejected, or an + Creates the output directory (with parents) if it does not + exist. Each output artifact is written as soon as its + content is computed, so when the input is rejected, or an extra keyword duplicates a group name or another extra keyword, the output directory holds whatever the steps before the failing one produced, and the error message names what @@ -719,8 +701,8 @@ def main(argv: list[str] | None = None) -> int: """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) - args.output_dir.mkdir(parents=True, exist_ok=True) try: + args.output_dir.mkdir(parents=True, exist_ok=True) source: PooledKeywords = KeywordPooler( args.run_dir_1, args.run_dir_2, args.output_dir).run() clusters: KeywordGroups = KeywordClusterer( @@ -735,7 +717,7 @@ def main(argv: list[str] | None = None) -> int: f"Done. Clustered {len(source.keywords)} keywords into" f" {len(clusters.names)}. {elapsed} elapsed.", file=sys.stderr) - except ClusterError as error: + except (ClusterError, OSError) as error: print(f"error: {error}", file=sys.stderr) return 1 return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/export_llm_input.py b/tools/src/pop_fem_audit_tools/commands/export_llm_input.py index 542c9fb..4b8464e 100644 --- a/tools/src/pop_fem_audit_tools/commands/export_llm_input.py +++ b/tools/src/pop_fem_audit_tools/commands/export_llm_input.py @@ -11,26 +11,17 @@ project's lyrics-only firewall: the output carries only the lyrics text of each song, identified by an opaque song key; no title, artist, or chart data crosses into the LLM input. -With ``--extras``, each record's ``content`` becomes a JSON -object serialized as a string, its ``lyrics`` key holding the -song's lyrics followed by the keys of the given extras file in -their file order, so a step that needs parameters alongside the -lyrics can carry them without this module knowing what they mean. - -With ``--extras-per-id``, the same merge happens per song: the -given file maps a song ID to the extra keys of that one song, and -the export is restricted to the song IDs the file names, so a step -that revisits only some of the songs, each with its own parameters, -gets exactly those records. The two options may be given together, -in which case a record's keys are ``lyrics``, the shared extras' -keys, then that song's own keys, each group in its file order. +With ``--extras`` and ``--extras-per-id``, a record's content may +carry extra parameters alongside the lyrics, so a step that needs +them can get them without this module knowing what they mean; see +the exporter's content-building step for how the two merge. """ import argparse import json import sys import time from pathlib import Path -from typing import Any +from typing import Any, ClassVar import sqlalchemy as sa from sqlalchemy.orm import Session @@ -40,6 +31,270 @@ from ..models import Song from ..utils import format_duration +class LlmInputExporter: + """The exporter of the LLM input JSONL file.""" + + __LYRICS_KEY: ClassVar[str] = "lyrics" + """The key holding the lyrics in a record's merged content, + and the key forbidden in an extras file.""" + + def __init__( + self, output_jsonl: Path, extras: Path | None = None, + extras_per_id: Path | None = None) -> None: + """Set up the exporter of the LLM input JSONL file. + + :param output_jsonl: The JSONL output file. + :param extras: The extras JSON file, or None for none. + :param extras_per_id: The per-ID extras JSON file, or + None for none. + """ + self.__output_jsonl: Path = output_jsonl + """The JSONL output file.""" + self.__extras_path: Path | None = extras + """The extras JSON file, or None for none.""" + self.__extras_per_id_path: Path | None = extras_per_id + """The per-ID extras JSON file, or None for none.""" + + def run(self) -> int: + """Export the songs' lyrics to the output JSONL file. + + :return: The number of songs exported. + :raises OSError: When a file cannot be read or written. + :raises sqlalchemy.exc.SQLAlchemyError: When the working + store cannot be read. + :raises ValueError: When an extras file is malformed, an + exported song has no lyrics, or the per-ID extras + name a song the working store does not have. + """ + session: Session = ds.get_db() + try: + extras: dict[str, Any] | None = None + if self.__extras_path is not None: + extras = self.__load_extras(self.__extras_path) + extras_per_id: dict[str, dict[str, Any]] | None = None + if self.__extras_per_id_path is not None: + extras_per_id = self.__load_extras_per_id( + self.__extras_per_id_path) + lines: list[str] = self.__build_lines( + session, extras, extras_per_id) + finally: + session.close() + self.__write_output(lines) + return len(lines) + + @staticmethod + def __no_duplicate_keys( + pairs: list[tuple[str, Any]]) -> dict[str, Any]: + """Build a dict from JSON object pairs, rejecting + duplicates. + + :param pairs: The key-value pairs of a JSON object, in + file order. + :return: The pairs as a dict, in file order. + :raises ValueError: When a key appears more than once. + """ + result: dict[str, Any] = {} + key: str + value: Any + for key, value in pairs: + if key in result: + raise ValueError( + f"duplicate key \"{key}\" in extras") + result[key] = value + return result + + @classmethod + def __load_json_object( + cls, path: Path, label: str) -> dict[str, Any]: + """Load a single JSON object from a file, in file order. + + :param path: The JSON file. + :param label: The kind of file, for the error messages. + :return: The object, in file order. + :raises OSError: When the file cannot be read. + :raises ValueError: When the file is not valid JSON, is + not a JSON object, or has duplicate keys. + """ + with open(path, encoding="utf-8") as file: + text: str = file.read() + try: + data: Any = json.loads( + text, object_pairs_hook=cls.__no_duplicate_keys) + except json.JSONDecodeError as error: + raise ValueError( + f"invalid JSON in {label} file {path}: {error}") \ + from error + if not isinstance(data, dict): + raise ValueError( + f"{label} file {path} must contain a JSON object") + return data + + @classmethod + def __load_extras(cls, path: Path) -> dict[str, Any]: + """Load the extras object from a JSON file. + + :param path: The extras JSON file. + :return: The extras, in file order. + :raises OSError: When the file cannot be read. + :raises ValueError: When the file is not valid JSON, is + not a JSON object, has duplicate keys, or has a + "lyrics" key. + """ + data: dict[str, Any] = cls.__load_json_object( + path, "extras") + if cls.__LYRICS_KEY in data: + raise ValueError( + f"extras file {path} must not have a" + f" \"{cls.__LYRICS_KEY}\" key") + return data + + @classmethod + def __load_extras_per_id( + cls, path: Path) -> dict[str, dict[str, Any]]: + """Load the per-ID extras object from a JSON file. + + :param path: The per-ID extras JSON file, mapping a song + ID, as ``song-``, to the extras of that one song. + :return: The extras of each song ID, in file order, every + song's own extras in their file order too. + :raises OSError: When the file cannot be read. + :raises ValueError: When the file is not valid JSON, is + not a JSON object, has duplicate keys, has a song + whose value is not a JSON object, or has a song with + a "lyrics" key. + """ + data: dict[str, Any] = cls.__load_json_object( + path, "per-ID extras") + song_id: str + extras: Any + for song_id, extras in data.items(): + if not isinstance(extras, dict): + raise ValueError( + f"per-ID extras file {path}: id {song_id}" + " must have a JSON object") + if cls.__LYRICS_KEY in extras: + raise ValueError( + f"per-ID extras file {path}: id {song_id}" + f" must not have a \"{cls.__LYRICS_KEY}\"" + " key") + return data + + def __build_lines( + self, session: Session, + extras: dict[str, Any] | None = None, + extras_per_id: dict[str, dict[str, Any]] | None + = None) -> list[str]: + """Build the JSONL lines of the exported songs' lyrics. + + Every song is exported, unless per-ID extras are given, + in which case only the songs they name are; see + :meth:`__build_content` for how the extras merge into a + record's content. + + :param session: The database session. + :param extras: The extra parameters merged into every + record's content alongside the lyrics, in the order + they are to appear, or None for none. + :param extras_per_id: The extra parameters merged into + the content of one record alone, keyed by that + record's song ID and in the order they are to appear, + restricting the export to the song IDs they name, or + None for no such extras and no such restriction. + :return: The JSON lines, one per exported song, ordered by + song ID. + :raises ValueError: When an exported song has no lyrics, + or the per-ID extras name a song the working store + does not have. + """ + lines: list[str] = [] + exported: set[str] = set() + song: Song + for song in session.scalars( + sa.select(Song).order_by(Song.id)): + song_id: str = f"song-{song.id}" + if extras_per_id is not None \ + and song_id not in extras_per_id: + continue + if song.lyrics is None: + raise ValueError( + f"song {song.id} \"{song.title}\": no lyrics") + song_extras: dict[str, Any] | None = None \ + if extras_per_id is None \ + else extras_per_id[song_id] + content: str = self.__build_content( + song.lyrics, extras, song_extras) + record: dict[str, str] = { + "id": song_id, "content": content} + lines.append(json.dumps(record, ensure_ascii=False)) + exported.add(song_id) + if extras_per_id is not None: + missing: list[str] = sorted( + set(extras_per_id) - exported) + if len(missing) > 0: + raise ValueError( + "the per-ID extras name songs the working" + f" store does not have: {', '.join(missing)}") + return lines + + @classmethod + def __build_content( + cls, lyrics: str, extras: dict[str, Any] | None, + song_extras: dict[str, Any] | None) -> str: + """Build the content of one exported record. + + Without extras of either kind, a record's content is the + bare lyrics string. With ``--extras``, the content + becomes a JSON object serialized as a string, its + "lyrics" key holding the song's lyrics followed by the + keys of the given extras file, in their file order. With + ``--extras-per-id``, the same merge happens per song: the + song's own extra keys follow the lyrics instead. When + both are given, a record's keys are "lyrics", the shared + extras' keys, then that song's own keys, each group in + its file order. + + :param lyrics: The lyrics of the song. + :param extras: The extra parameters shared by every + record, in the order they are to appear, or None for + none. + :param song_extras: The extra parameters of this record + alone, in the order they are to appear, or None for + none. + :return: The bare lyrics when there are no extras of + either kind, or otherwise a JSON object serialized as + a string, whose first key is "lyrics" holding the + lyrics, followed by the shared extras' keys and then + this record's own keys, each group in its given + order. + """ + if extras is None and song_extras is None: + return lyrics + payload: dict[str, Any] = {cls.__LYRICS_KEY: lyrics} + if extras is not None: + payload.update(extras) + if song_extras is not None: + payload.update(song_extras) + return json.dumps(payload, ensure_ascii=False) + + def __write_output(self, lines: list[str]) -> None: + """Write the exported lines to the output JSONL file. + + Creates the parent directory when it does not exist. + + :param lines: The JSONL lines, in the output order. + :return: None. + :raises OSError: When the file cannot be written. + """ + self.__output_jsonl.parent.mkdir( + parents=True, exist_ok=True) + with open( + self.__output_jsonl, "w", + encoding="utf-8") as file: + line: str + for line in lines: + file.write(line + "\n") + + def parse_args(argv: list[str] | None) -> argparse.Namespace: """Parse the command-line arguments. @@ -56,227 +311,33 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: parser.add_argument( "--extras", type=Path, default=None, help="a JSON file holding a single JSON object of extra" - " parameters; when given, each record's \"content\"" - " becomes a JSON object string with a \"lyrics\" key" - " followed by the extras' keys, instead of the bare" - " lyrics string") + " parameters merged into every record's content") parser.add_argument( "--extras-per-id", type=Path, default=None, help="a JSON file holding a single JSON object that maps a" " song ID, as \"song-\", to a JSON object of extra" - " parameters for that one song; the song's object is" - " merged into its \"content\" the same way as with" - " --extras, and the export is restricted to the song" - " IDs the file names") + " parameters for that one song, restricting the" + " export to the song IDs the file names") return parser.parse_args(argv) -def __no_duplicate_keys( - pairs: list[tuple[str, Any]]) -> dict[str, Any]: - """Build a dict from JSON object pairs, rejecting duplicates. - - :param pairs: The key-value pairs of a JSON object, in file - order. - :return: The pairs as a dict, in file order. - :raises ValueError: When a key appears more than once. - """ - result: dict[str, Any] = {} - key: str - value: Any - for key, value in pairs: - if key in result: - raise ValueError(f"duplicate key \"{key}\" in extras") - result[key] = value - return result - - -def __load_json_object(path: Path, label: str) -> dict[str, Any]: - """Load a single JSON object from a file, in file order. - - :param path: The JSON file. - :param label: The kind of file, for the error messages. - :return: The object, in file order. - :raises OSError: When the file cannot be read. - :raises ValueError: When the file is not valid JSON, is not - a JSON object, or has duplicate keys. - """ - with open(path, encoding="utf-8") as file: - text: str = file.read() - try: - data: Any = json.loads( - text, object_pairs_hook=__no_duplicate_keys) - except json.JSONDecodeError as error: - raise ValueError( - f"invalid JSON in {label} file {path}: {error}") \ - from error - if not isinstance(data, dict): - raise ValueError( - f"{label} file {path} must contain a JSON object") - return data - - -def load_extras(path: Path) -> dict[str, Any]: - """Load the extras object from a JSON file. - - :param path: The extras JSON file. - :return: The extras, in file order. - :raises OSError: When the file cannot be read. - :raises ValueError: When the file is not valid JSON, is not - a JSON object, has duplicate keys, or has a "lyrics" key. - """ - data: dict[str, Any] = __load_json_object(path, "extras") - if "lyrics" in data: - raise ValueError( - f"extras file {path} must not have a \"lyrics\" key") - return data - - -def load_extras_per_id(path: Path) -> dict[str, dict[str, Any]]: - """Load the per-ID extras object from a JSON file. - - :param path: The per-ID extras JSON file, mapping a song ID, - as ``song-``, to the extras of that one song. - :return: The extras of each song ID, in file order, every - song's own extras in their file order too. - :raises OSError: When the file cannot be read. - :raises ValueError: When the file is not valid JSON, is not - a JSON object, has duplicate keys, has a song whose value - is not a JSON object, or has a song with a "lyrics" key. - """ - data: dict[str, Any] = __load_json_object(path, "per-ID extras") - song_id: str - extras: Any - for song_id, extras in data.items(): - if not isinstance(extras, dict): - raise ValueError( - f"per-ID extras file {path}: id {song_id} must" - " have a JSON object") - if "lyrics" in extras: - raise ValueError( - f"per-ID extras file {path}: id {song_id} must not" - " have a \"lyrics\" key") - return data - - -def __build_content( - lyrics: str, - extras: dict[str, Any] | None, - song_extras: dict[str, Any] | None) -> str: - """Build the content of one exported record. - - :param lyrics: The lyrics of the song. - :param extras: The extra parameters shared by every record, - in the order they are to appear, or None for none. - :param song_extras: The extra parameters of this record - alone, in the order they are to appear, or None for none. - :return: The bare lyrics when there are no extras of either - kind, or otherwise a JSON object serialized as a string, - whose first key is ``"lyrics"`` holding the lyrics, - followed by the shared extras' keys and then this - record's own keys, each group in its given order. - """ - if extras is None and song_extras is None: - return lyrics - payload: dict[str, Any] = {"lyrics": lyrics} - if extras is not None: - payload.update(extras) - if song_extras is not None: - payload.update(song_extras) - return json.dumps(payload, ensure_ascii=False) - - -def build_lines( - session: Session, - extras: dict[str, Any] | None = None, - extras_per_id: dict[str, dict[str, Any]] | None = None) \ - -> list[str]: - """Build the JSONL lines of the exported songs' lyrics. - - Without extras of either kind, each record's ``content`` is - the bare lyrics string. With extras, ``content`` is a JSON - object serialized as a string, whose first key is ``"lyrics"`` - holding the lyrics string, followed by the shared extras' keys - and then the song's own per-ID extras' keys, each group in its - given order. - - Every song is exported, unless per-ID extras are given, in - which case only the songs they name are. - - :param session: The database session. - :param extras: The extra parameters merged into every - record's content alongside the lyrics, in the order they - are to appear, or None for none. - :param extras_per_id: The extra parameters merged into the - content of one record alone, keyed by that record's song - ID and in the order they are to appear, restricting the - export to the song IDs they name, or None for no such - extras and no such restriction. - :return: The JSON lines, one per exported song, ordered by - song ID. - :raises ValueError: When an exported song has no lyrics, or - the per-ID extras name a song the working store does not - have. - """ - lines: list[str] = [] - exported: set[str] = set() - song: Song - for song in session.scalars(sa.select(Song).order_by(Song.id)): - song_id: str = f"song-{song.id}" - if extras_per_id is not None and song_id not in extras_per_id: - continue - if song.lyrics is None: - raise ValueError( - f"song {song.id} \"{song.title}\": no lyrics") - song_extras: dict[str, Any] | None = None \ - if extras_per_id is None else extras_per_id[song_id] - content: str = __build_content( - song.lyrics, extras, song_extras) - record: dict[str, str] = { - "id": song_id, "content": content} - lines.append(json.dumps(record, ensure_ascii=False)) - exported.add(song_id) - if extras_per_id is not None: - missing: list[str] = sorted(set(extras_per_id) - exported) - if len(missing) > 0: - raise ValueError( - "the per-ID extras name songs the working store" - f" does not have: {', '.join(missing)}") - return lines - - def main(argv: list[str] | None = None) -> int: """Export the LLM input JSONL file from the working store. - Every song is exported, unless ``--extras-per-id`` is given, - in which case only the songs its file names are. - :param argv: The command-line arguments, or None for ``sys.argv``. :return: The exit status: 0 on success, non-zero on failure. """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) - session: Session = ds.get_db() - lines: list[str] try: - extras: dict[str, Any] | None = None - if args.extras is not None: - extras = load_extras(args.extras) - extras_per_id: dict[str, dict[str, Any]] | None = None - if args.extras_per_id is not None: - extras_per_id = load_extras_per_id(args.extras_per_id) - lines = build_lines(session, extras, extras_per_id) + count: int = LlmInputExporter( + args.output_jsonl, args.extras, + args.extras_per_id).run() except (OSError, sa.exc.SQLAlchemyError, ValueError) as error: print(f"error: {error}", file=sys.stderr) return 1 - finally: - session.close() - args.output_jsonl.parent.mkdir(parents=True, exist_ok=True) - with open(args.output_jsonl, "w", encoding="utf-8") as file: - line: str - for line in lines: - file.write(line + "\n") elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. {len(lines)} songs exported." + print(f"Done. {count} songs exported." f" {elapsed} elapsed.", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/fetch_artists.py b/tools/src/pop_fem_audit_tools/commands/fetch_artists.py index 1c6b1e3..68140ad 100644 --- a/tools/src/pop_fem_audit_tools/commands/fetch_artists.py +++ b/tools/src/pop_fem_audit_tools/commands/fetch_artists.py @@ -12,12 +12,10 @@ layer. The working store is only read, never written; the ``build-db`` subcommand assembles the captured files into the store on the next rebuild. -Every fetched row is meant for later human verification: the -description of the resolved item is recorded in the note column -so that a bad match can be spotted. An unresolved artist or an -error on one artist is noted on its row and does not fail the -run. A row whose name is no longer an artist of the store is -dropped from the snapshot and reported on the standard error. +An unresolved artist or an error on one artist is noted on its +row and does not fail the run. A row whose name is no longer an +artist of the store is dropped from the snapshot and reported on +the standard error. """ import argparse import csv @@ -33,7 +31,7 @@ import urllib.request from collections.abc import Container, Sequence from dataclasses import asdict, dataclass, field, fields from pathlib import Path -from typing import Any, Literal, TextIO +from typing import Any, ClassVar, Literal, TextIO import sqlalchemy as sa from sqlalchemy.orm import Session @@ -43,71 +41,6 @@ from ..database import ds from ..models import Artist, Song, SongArtist from ..utils import format_duration -API_URL: str = "https://www.wikidata.org/w/api.php" -"""The URL of the Wikidata API endpoint.""" -SPARQL_URL: str = "https://query.wikidata.org/sparql" -"""The URL of the Wikidata Query Service SPARQL endpoint.""" -USER_AGENT: str = ( - f"pop-fem-audit-tools/{VERSION}" - " (https://github.com/imacat/pop-fem-audit;" - " mailto:imacat@mail.imacat.idv.tw)") -"""The User-Agent header sent on every HTTP request.""" -TIMEOUT: float = 30.0 -"""The timeout of an API HTTP request, in seconds.""" -SPARQL_TIMEOUT: float = 90.0 -"""The timeout of a SPARQL HTTP request, in seconds. - -Higher than the API timeout: the WDQS server aborts a slow -query at 60 seconds, and a lower client timeout would race -that server-side abort and misclassify a slow-but-answerable -query as a client-side timeout instead of letting the server's -own HTTP error response arrive and enter the retry path.""" -SLEEP_SECONDS: float = 1.0 -"""The delay between consecutive HTTP requests, in seconds.""" -MAX_ATTEMPTS: int = 5 -"""The maximum number of attempts on a transient error.""" -RETRY_SECONDS: float = 15.0 -"""The back-off unit on a transient error, in seconds; -multiplied by the attempt number already made.""" -RETRY_STATUSES: frozenset[int] = frozenset({429, 500, 502, 503}) -"""The HTTP statuses that are retried with a back-off.""" -MAX_STAGE1_TITLES: int = 3 -"""The maximum number of charted titles used for the stage-1 song -corroboration.""" -HUMAN_QID: str = "Q5" -"""The Wikidata item ID of "human".""" -ENSEMBLE_QID: str = "Q2088357" -"""The Wikidata item ID of "musical ensemble".""" -ORIGINAL_CAST_QID: str = "Q106497009" -"""The Wikidata item ID of "original cast".""" -GROUP_KEYWORDS: Sequence[str] = ("band", "group", "duo", "trio") -"""The label keywords that suggest a musical ensemble, covering -labels like "boy band" and "girl group".""" -NOTE_NOT_FOUND: str = "not found" -"""The note sentinel of an artist without a resolved Wikidata -item, written to the snapshot and read back for the -classification.""" -CORPUS_START_YEAR: int = 2016 -"""The first year of the corpus window: a member who left a group -before it never performed a corpus song.""" -MIXED_GENDER: str = "mixed" -"""The gender recorded for a group whose members do not share one -gender.""" -TIME_YEAR_PATTERN: re.Pattern[str] = re.compile(r"^[+-]?\d+") -"""The leading year of a Wikidata time value.""" -PINNED_QIDS: dict[str, str] = {} -"""The last-resort pinned item IDs, keyed by the artist name. - -An entry is for an artist the algorithm documented on -``ArtistFetcher`` is structurally unable to resolve, with its -justification recorded here. Currently empty: the only pin ever -needed, "Pinkfong" (typed as a brand, which the type gate -excludes by design), became moot when the store's artist entity -behind that credit was identified as Hope Segoine. - -A pinned name skips the candidate retrieval and corroboration -steps; its item ID is used directly.""" - class ArtistType(enum.StrEnum): """The decided artist type of a snapshot row.""" @@ -147,15 +80,14 @@ class ArtistSnapshot: return asdict(self) -SNAPSHOT_FIELDS: Sequence[str] = tuple( - x.name for x in fields(ArtistSnapshot)) -"""The header columns of the Wikidata artist snapshot CSV file.""" - - @dataclass class GroupMember: """One has-part member of a Wikidata group item.""" + __CORPUS_START_YEAR: ClassVar[int] = 2016 + """The first year of the corpus window: a member who left a + group before it never performed a corpus song.""" + qid: str """The item ID of the member.""" start_years: list[int] = field(default_factory=list) @@ -181,7 +113,7 @@ class GroupMember: if len(self.end_years) == 0: return True last_end: int = max(self.end_years) - if last_end >= CORPUS_START_YEAR: + if last_end >= self.__CORPUS_START_YEAR: return True if len(self.start_years) == 0: return False @@ -229,22 +161,6 @@ class RetryExhausted(Exception): """ -def parse_args(argv: list[str] | None) -> argparse.Namespace: - """Parse the command-line arguments. - - :param argv: The command-line arguments, or None for - ``sys.argv``. - :return: The parsed arguments. - """ - parser: argparse.ArgumentParser = argparse.ArgumentParser( - description="Fetch the artist metadata from Wikidata" - " into the capture layer.") - parser.add_argument( - "wikidata_csv", type=Path, - help="the Wikidata artist snapshot CSV file") - return parser.parse_args(argv) - - class ArtistFetcher: """A fetcher of artist metadata from Wikidata. @@ -291,6 +207,74 @@ class ArtistFetcher: in the note. """ + __API_URL: ClassVar[str] = "https://www.wikidata.org/w/api.php" + """The URL of the Wikidata API endpoint.""" + __SPARQL_URL: ClassVar[str] \ + = "https://query.wikidata.org/sparql" + """The URL of the Wikidata Query Service SPARQL endpoint.""" + __USER_AGENT: ClassVar[str] = ( + f"pop-fem-audit-tools/{VERSION}" + " (https://github.com/imacat/pop-fem-audit;" + " mailto:imacat@mail.imacat.idv.tw)") + """The User-Agent header sent on every HTTP request.""" + __TIMEOUT: ClassVar[float] = 30.0 + """The timeout of an API HTTP request, in seconds.""" + __SPARQL_TIMEOUT: ClassVar[float] = 90.0 + """The timeout of a SPARQL HTTP request, in seconds. + + Higher than the API timeout: the WDQS server aborts a slow + query at 60 seconds, and a lower client timeout would race + that server-side abort and misclassify a slow-but-answerable + query as a client-side timeout instead of letting the + server's own HTTP error response arrive and enter the retry + path.""" + __SLEEP_SECONDS: ClassVar[float] = 1.0 + """The delay between consecutive HTTP requests, in + seconds.""" + __MAX_ATTEMPTS: ClassVar[int] = 5 + """The maximum number of attempts on a transient error.""" + __RETRY_SECONDS: ClassVar[float] = 15.0 + """The back-off unit on a transient error, in seconds; + multiplied by the attempt number already made.""" + __RETRY_STATUSES: ClassVar[frozenset[int]] \ + = frozenset({429, 500, 502, 503}) + """The HTTP statuses that are retried with a back-off.""" + __MAX_STAGE1_TITLES: ClassVar[int] = 3 + """The maximum number of charted titles used for the + stage-1 song corroboration.""" + __HUMAN_QID: ClassVar[str] = "Q5" + """The Wikidata item ID of "human".""" + __ENSEMBLE_QID: ClassVar[str] = "Q2088357" + """The Wikidata item ID of "musical ensemble".""" + __ORIGINAL_CAST_QID: ClassVar[str] = "Q106497009" + """The Wikidata item ID of "original cast".""" + __GROUP_KEYWORDS: ClassVar[Sequence[str]] \ + = ("band", "group", "duo", "trio") + """The label keywords that suggest a musical ensemble, + covering labels like "boy band" and "girl group".""" + __NOTE_NOT_FOUND: ClassVar[str] = "not found" + """The note sentinel of an artist without a resolved + Wikidata item.""" + __MIXED_GENDER: ClassVar[str] = "mixed" + """The gender recorded for a group whose members do not + share one gender.""" + __TIME_YEAR_PATTERN: ClassVar[re.Pattern[str]] \ + = re.compile(r"^[+-]?\d+") + """The leading year of a Wikidata time value.""" + __PINNED_QIDS: ClassVar[dict[str, str]] = {} + """The last-resort pinned item IDs, keyed by the artist name. + + An entry is for an artist the algorithm documented on + ``ArtistFetcher`` is structurally unable to resolve, with its + justification recorded here. Currently empty: the only pin + ever needed, "Pinkfong" (typed as a brand, which the type + gate excludes by design), became moot when the store's + artist entity behind that credit was identified as Hope + Segoine. + + A pinned name skips the candidate retrieval and corroboration + steps; its item ID is used directly.""" + def __init__(self) -> None: """Construct the fetcher.""" self.__sent: int = 0 @@ -317,7 +301,7 @@ class ArtistFetcher: try: qid: str | None = self.__resolve_qid(name, titles) if qid is None: - snapshot.note = NOTE_NOT_FOUND + snapshot.note = self.__NOTE_NOT_FOUND return snapshot snapshot.qid = qid self.__resolve(snapshot) @@ -342,8 +326,8 @@ class ArtistFetcher: transient error are exhausted. :raises ValueError: On a JSON decoding error. """ - if name in PINNED_QIDS: - return PINNED_QIDS[name] + if name in self.__PINNED_QIDS: + return self.__PINNED_QIDS[name] candidates: list[str] = self.__candidates(name) if len(candidates) == 0: return None @@ -373,11 +357,11 @@ class ArtistFetcher: {{ ?item rdfs:label ?name }} UNION {{ ?item skos:altLabel ?name }} {{ - ?item wdt:P31 wd:{HUMAN_QID} + ?item wdt:P31 wd:{self.__HUMAN_QID} }} UNION {{ - ?item wdt:P31/wdt:P279* wd:{ENSEMBLE_QID} + ?item wdt:P31/wdt:P279* wd:{self.__ENSEMBLE_QID} }} UNION {{ - ?item wdt:P31 wd:{ORIGINAL_CAST_QID} + ?item wdt:P31 wd:{self.__ORIGINAL_CAST_QID} }} }} """ @@ -403,7 +387,7 @@ class ArtistFetcher: transient error are exhausted. :raises ValueError: On a JSON decoding error. """ - subset: Sequence[str] = titles[:MAX_STAGE1_TITLES] + subset: Sequence[str] = titles[:self.__MAX_STAGE1_TITLES] if len(subset) == 0: return None query: str = f""" @@ -532,7 +516,7 @@ class ArtistFetcher: qid: str for qid in qids: member: MemberClaims = claims.get(qid, MemberClaims()) - if HUMAN_QID not in member.instance_of_ids: + if self.__HUMAN_QID not in member.instance_of_ids: continue if len(member.gender_ids) == 0: return @@ -543,7 +527,7 @@ class ArtistFetcher: [x[1] for x in genders], any_language=True) unique: set[str] = {x[1] for x in genders} snapshot.gender = labels[genders[0][1]] \ - if len(unique) == 1 else MIXED_GENDER + if len(unique) == 1 else self.__MIXED_GENDER basis: str = "gender derived from members: " + "; ".join( f"{x} {labels[y]}" for x, y in genders) snapshot.note = f"{snapshot.note}; {basis}" \ @@ -677,7 +661,8 @@ class ArtistFetcher: or not isinstance(value.get("time"), str): continue match: re.Match[str] | None \ - = TIME_YEAR_PATTERN.match(value["time"]) + = ArtistFetcher.__TIME_YEAR_PATTERN.match( + value["time"]) if match is not None: years.append(int(match.group())) return years @@ -806,12 +791,13 @@ class ArtistFetcher: ``ArtistType.GROUP`` for a musical ensemble, or the empty string for the human to decide. """ - if HUMAN_QID in type_ids: + if ArtistFetcher.__HUMAN_QID in type_ids: return ArtistType.SOLO qid: str for qid in type_ids: label: str = labels.get(qid, "").lower() - if any(x in label for x in GROUP_KEYWORDS): + if any(x in label + for x in ArtistFetcher.__GROUP_KEYWORDS): return ArtistType.GROUP return "" @@ -827,14 +813,15 @@ class ArtistFetcher: transient error are exhausted. :raises ValueError: On a JSON decoding error. """ - url: str = (f"{SPARQL_URL}?" - f"{urllib.parse.urlencode({'query': query})}") + url: str = ( + f"{self.__SPARQL_URL}?" + f"{urllib.parse.urlencode({'query': query})}") request: urllib.request.Request = urllib.request.Request( url, headers={ - "User-Agent": USER_AGENT, + "User-Agent": self.__USER_AGENT, "Accept": "application/sparql-results+json"}) body: bytes = self.__send( - request, timeout=SPARQL_TIMEOUT) + request, timeout=self.__SPARQL_TIMEOUT) data: Any = json.loads(body) bindings: Any = None if isinstance(data, dict) \ @@ -868,13 +855,14 @@ class ArtistFetcher: transient error are exhausted. :raises ValueError: On a JSON decoding error. """ - url: str = f"{API_URL}?{urllib.parse.urlencode(params)}" + url: str \ + = f"{self.__API_URL}?{urllib.parse.urlencode(params)}" request: urllib.request.Request = urllib.request.Request( - url, headers={"User-Agent": USER_AGENT}) + url, headers={"User-Agent": self.__USER_AGENT}) return json.loads(self.__send(request)) def __send(self, request: urllib.request.Request, - timeout: float = TIMEOUT) -> bytes: + timeout: float = __TIMEOUT) -> bytes: """Send an HTTP request, retrying on a transient error. Consecutive requests are separated by a fixed delay. A @@ -892,7 +880,7 @@ class ArtistFetcher: transient error are exhausted. """ if self.__sent > 0: - time.sleep(SLEEP_SECONDS) + time.sleep(self.__SLEEP_SECONDS) self.__sent += 1 attempt: int = 1 reason: str | None @@ -905,10 +893,10 @@ class ArtistFetcher: reason = self.__retry_reason(error) if reason is None: raise - if attempt >= MAX_ATTEMPTS: + if attempt >= self.__MAX_ATTEMPTS: raise RetryExhausted( f"retries exhausted ({reason})") from error - time.sleep(RETRY_SECONDS * attempt) + time.sleep(self.__RETRY_SECONDS * attempt) attempt += 1 @staticmethod @@ -923,7 +911,7 @@ class ArtistFetcher: the error is not transient and must not be retried. """ if isinstance(error, urllib.error.HTTPError): - if error.code not in RETRY_STATUSES: + if error.code not in ArtistFetcher.__RETRY_STATUSES: return None return str(error) if isinstance(error, TimeoutError): @@ -968,92 +956,216 @@ class ArtistFetcher: return uri.rsplit("/", 1)[-1] -def read_snapshot_rows(file: TextIO) -> list[dict[str, str]]: - """Read the current rows of a snapshot CSV file handle. +@dataclass(frozen=True) +class FetchCounts: + """The outcome counts of one snapshot update run.""" - :param file: The open, seekable snapshot CSV file. - :return: The rows, keyed by the column name. - :raises OSError: When the file cannot be read. + fetched: int + """The number of artists newly resolved.""" + not_found: int + """The number of artists left unresolved.""" + errors: int + """The number of artists that ended in an error.""" + + +class ArtistSnapshotUpdater: + """The updater of the Wikidata artist snapshot CSV file. + + Fetches the metadata of every artist of the working store + that the snapshot does not resolve yet, appends a row for + each to the snapshot as it is fetched, and rewrites the + snapshot sorted by artist name with its stale rows dropped. """ - file.seek(0) - reader: csv.DictReader[str] = csv.DictReader(file) - return list(reader) + __SNAPSHOT_FIELDS: ClassVar[Sequence[str]] = tuple( + x.name for x in fields(ArtistSnapshot)) + """The header columns of the Wikidata artist snapshot CSV file.""" -def read_artist_titles(session: Session, - artist_id: int) -> list[str]: - """Read the charted song titles credited to an artist. + def __init__(self, wikidata_csv: Path) -> None: + """Set up the updater. - :param session: The database session. - :param artist_id: The artist ID. - :return: The song titles credited to the artist, ordered by - the song ID, with the duplicate titles removed. - """ - titles: Sequence[str] = session.scalars( - sa.select(Song.title) - .join(SongArtist, SongArtist.song_id == Song.id) - .where(SongArtist.artist_id == artist_id) - .order_by(Song.id)).all() - return list(dict.fromkeys(titles)) + :param wikidata_csv: The Wikidata artist snapshot CSV + file. + """ + self.__wikidata_csv: Path = wikidata_csv + """The Wikidata artist snapshot CSV file.""" + def run(self) -> FetchCounts: + """Fetch every unresolved artist and update the snapshot. -def ensure_snapshot_header(file: TextIO) -> None: - """Write the snapshot CSV header row if the file is empty. + :return: The counts of the run. + :raises OSError: When the snapshot file, or its parent + directory, cannot be read or written. + :raises sqlalchemy.exc.SQLAlchemyError: When the working + store cannot be read. + """ + session: Session = ds.get_db() + try: + return self.__run(session) + finally: + session.close() - :param file: The open, seekable snapshot CSV file. - :return: None. - :raises OSError: When the file cannot be written. - """ - file.seek(0, os.SEEK_END) - if file.tell() == 0: - csv.writer(file).writerow(SNAPSHOT_FIELDS) + def __run(self, session: Session) -> FetchCounts: + """Run the fetch loop with an open database session. + + :param session: The database session. + :return: The counts of the run. + :raises OSError: When the snapshot file, or its parent + directory, cannot be read or written. + """ + fetcher: ArtistFetcher = ArtistFetcher() + fetched: int = 0 + not_found: int = 0 + errors: int = 0 + self.__wikidata_csv.parent.mkdir( + parents=True, exist_ok=True) + with open(self.__wikidata_csv, "a+", encoding="utf-8", + newline="") as csv_file: + done: set[str] = { + x["name"] for x in + self.__read_snapshot_rows(csv_file) + if x["gender"] != ""} + self.__ensure_snapshot_header(csv_file) + names: set[str] = set() + artist: Artist + for artist in session.scalars( + sa.select(Artist).order_by(Artist.id)): + names.add(artist.name) + if artist.name in done: + continue + titles: list[str] = self.__read_artist_titles( + session, artist.id) + snapshot: ArtistSnapshot = fetcher.fetch( + artist.name, titles) + self.__append_row(csv_file, snapshot) + status: str = snapshot.qid + if snapshot.note == "not found": + not_found += 1 + status = "not found" + elif snapshot.note.startswith("error: "): + errors += 1 + status = snapshot.note + else: + fetched += 1 + print(f"artist \"{artist.name}\": {status}", + file=sys.stderr) + self.__write_snapshot(csv_file, names) + return FetchCounts( + fetched=fetched, not_found=not_found, errors=errors) + + @staticmethod + def __read_snapshot_rows(file: TextIO) \ + -> list[dict[str, str]]: + """Read the current rows of a snapshot CSV file handle. + + :param file: The open, seekable snapshot CSV file. + :return: The rows, keyed by the column name. + :raises OSError: When the file cannot be read. + """ + file.seek(0) + reader: csv.DictReader[str] = csv.DictReader(file) + return list(reader) + + @staticmethod + def __read_artist_titles(session: Session, + artist_id: int) -> list[str]: + """Read the charted song titles credited to an artist. + + :param session: The database session. + :param artist_id: The artist ID. + :return: The song titles credited to the artist, ordered + by the song ID, with the duplicate titles removed. + """ + titles: Sequence[str] = session.scalars( + sa.select(Song.title) + .join(SongArtist, SongArtist.song_id == Song.id) + .where(SongArtist.artist_id == artist_id) + .order_by(Song.id)).all() + return list(dict.fromkeys(titles)) + + @staticmethod + def __ensure_snapshot_header(file: TextIO) -> None: + """Write the snapshot CSV header row if the file is + empty. + + :param file: The open, seekable snapshot CSV file. + :return: None. + :raises OSError: When the file cannot be written. + """ + file.seek(0, os.SEEK_END) + if file.tell() == 0: + csv.writer(file).writerow( + ArtistSnapshotUpdater.__SNAPSHOT_FIELDS) + file.flush() + + @staticmethod + def __append_row(file: TextIO, + snapshot: ArtistSnapshot) -> None: + """Append a snapshot row to a snapshot CSV file handle. + + :param file: The open snapshot CSV file, opened for + append. + :param snapshot: The snapshot of an artist. + :return: None. + :raises OSError: When the file cannot be written. + """ + csv.DictWriter( + file, ArtistSnapshotUpdater.__SNAPSHOT_FIELDS).writerow( + snapshot.to_row()) file.flush() + @staticmethod + def __write_snapshot(file: TextIO, names: Container[str]) \ + -> None: + """Rewrite a snapshot CSV file handle sorted by artist + name. -def append_row(file: TextIO, snapshot: ArtistSnapshot) -> None: - """Append a snapshot row to a snapshot CSV file handle. + The rows are ordered by the case-folded artist name, + matching the convention of the derived ``artists.csv``. + An artist keeps one row only, the last one of the file, + so that a re-fetched artist replaces its earlier row. A + row whose name is not an artist of the store is dropped + and reported on the standard error. - :param file: The open snapshot CSV file, opened for append. - :param snapshot: The snapshot of an artist. - :return: None. - :raises OSError: When the file cannot be written. + :param file: The open, seekable snapshot CSV file. + :param names: The artist names of the working store. + :return: None. + :raises OSError: When the file cannot be read or written. + """ + kept: dict[str, dict[str, str]] = {} + row: dict[str, str] + for row in ArtistSnapshotUpdater.__read_snapshot_rows( + file): + if row["name"] not in names: + print(f"dropped stale row \"{row['name']}\":" + " no such artist in the store", + file=sys.stderr) + continue + kept[row["name"]] = row + ordered: list[dict[str, str]] = sorted( + kept.values(), key=lambda x: x["name"].casefold()) + file.seek(0) + file.truncate() + writer: csv.DictWriter[str] = csv.DictWriter( + file, ArtistSnapshotUpdater.__SNAPSHOT_FIELDS) + writer.writeheader() + writer.writerows(ordered) + + +def parse_args(argv: list[str] | None) -> argparse.Namespace: + """Parse the command-line arguments. + + :param argv: The command-line arguments, or None for + ``sys.argv``. + :return: The parsed arguments. """ - csv.DictWriter(file, SNAPSHOT_FIELDS).writerow( - snapshot.to_row()) - file.flush() - - -def write_snapshot(file: TextIO, names: Container[str]) -> None: - """Rewrite a snapshot CSV file handle sorted by artist name. - - The rows are ordered by the case-folded artist name, matching - the convention of the derived ``artists.csv``. An artist - keeps one row only, the last one of the file, so that a - re-fetched artist replaces its earlier row. A row whose name - is not an artist of the store is dropped and reported on the - standard error. - - :param file: The open, seekable snapshot CSV file. - :param names: The artist names of the working store. - :return: None. - :raises OSError: When the file cannot be read or written. - """ - kept: dict[str, dict[str, str]] = {} - row: dict[str, str] - for row in read_snapshot_rows(file): - if row["name"] not in names: - print(f"dropped stale row \"{row['name']}\":" - " no such artist in the store", file=sys.stderr) - continue - kept[row["name"]] = row - ordered: list[dict[str, str]] = sorted( - kept.values(), key=lambda x: x["name"].casefold()) - file.seek(0) - file.truncate() - writer: csv.DictWriter[str] = csv.DictWriter( - file, SNAPSHOT_FIELDS) - writer.writeheader() - writer.writerows(ordered) + parser: argparse.ArgumentParser = argparse.ArgumentParser( + description="Fetch the artist metadata from Wikidata" + " into the capture layer.") + parser.add_argument( + "wikidata_csv", type=Path, + help="the Wikidata artist snapshot CSV file") + return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: @@ -1066,52 +1178,16 @@ def main(argv: list[str] | None = None) -> int: """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) - fetcher: ArtistFetcher = ArtistFetcher() - fetched: int = 0 - not_found: int = 0 - errors: int = 0 - session: Session = ds.get_db() try: - args.wikidata_csv.parent.mkdir( - parents=True, exist_ok=True) - with open(args.wikidata_csv, "a+", encoding="utf-8", - newline="") as csv_file: - done: set[str] = {x["name"] for x in - read_snapshot_rows(csv_file) - if x["gender"] != ""} - ensure_snapshot_header(csv_file) - names: set[str] = set() - artist: Artist - for artist in session.scalars( - sa.select(Artist).order_by(Artist.id)): - names.add(artist.name) - if artist.name in done: - continue - titles: list[str] = read_artist_titles( - session, artist.id) - snapshot: ArtistSnapshot = fetcher.fetch( - artist.name, titles) - append_row(csv_file, snapshot) - status: str = snapshot.qid - if snapshot.note == NOTE_NOT_FOUND: - not_found += 1 - status = "not found" - elif snapshot.note.startswith("error: "): - errors += 1 - status = snapshot.note - else: - fetched += 1 - print(f"artist \"{artist.name}\": {status}", - file=sys.stderr) - write_snapshot(csv_file, names) + counts: FetchCounts \ + = ArtistSnapshotUpdater(args.wikidata_csv).run() except (OSError, sa.exc.SQLAlchemyError) as error: print(f"error: {error}", file=sys.stderr) return 1 - finally: - session.close() - attempted: int = fetched + not_found + errors + attempted: int = counts.fetched + counts.not_found \ + + counts.errors elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. Resolved {fetched}/{attempted} artists." - f" {elapsed} elapsed.", + print(f"Done. Resolved {counts.fetched}/{attempted}" + f" artists. {elapsed} elapsed.", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/fetch_lyrics.py b/tools/src/pop_fem_audit_tools/commands/fetch_lyrics.py index d801f6c..162521e 100644 --- a/tools/src/pop_fem_audit_tools/commands/fetch_lyrics.py +++ b/tools/src/pop_fem_audit_tools/commands/fetch_lyrics.py @@ -30,8 +30,9 @@ import time import urllib.parse import urllib.request from collections.abc import Sequence +from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, ClassVar import sqlalchemy as sa from sqlalchemy.orm import Session @@ -45,60 +46,6 @@ from ..models import ( ) from ..utils import format_duration -PROVENANCE_FIELDS: Sequence[str] = ( - "song_id", "source", "method", "acquired_at", "note") -"""The header columns of the lyrics provenance CSV file.""" -USER_AGENT: str = ("pop-fem-audit-tools" - " (https://github.com/imacat/pop-fem-audit)") -"""The User-Agent header sent on every HTTP request.""" -TIMEOUT: float = 30.0 -"""The timeout of an HTTP request, in seconds.""" -SLEEP_SECONDS: float = 1.0 -"""The delay between consecutive HTTP requests, in seconds.""" - - -def __build_normalization() -> dict[int, str | None]: - """Build the lyrics normalization translation table. - - :return: The codepoint-to-replacement mapping, a replacement - of None meaning removal. - """ - table: dict[int, str | None] = {} - codepoint: int - for codepoint in range(0x80, 0xa0): - try: - table[codepoint] = bytes([codepoint]).decode("cp1252") - except UnicodeDecodeError: - table[codepoint] = None - table[0x0435] = "e" - table[0x03cc] = "ó" - for codepoint in (0x2005, 0x205f, 0x200a): - table[codepoint] = " " - for codepoint in (0x200b, 0x200c, 0x200d, 0xfeff): - table[codepoint] = None - return table - - -NORMALIZATION: dict[int, str | None] = __build_normalization() -"""The codepoint-to-replacement mapping applied to fetched -lyrics: cp1252-mojibake restoration for U+0080-U+009F (with the -five byte values undefined in cp1252 removed), homoglyph -restoration for the Cyrillic "e" and the Greek "o" with tonos, -ASCII-space restoration for exotic space variants, and removal -of zero-width characters. A replacement of None removes the -codepoint.""" - - -def normalize_lyrics(text: str) -> str: - """Restore or remove watermark and mojibake characters. - - :param text: The lyrics text as fetched from an API. - :return: The text with the codepoints in - :data:`NORMALIZATION` replaced or removed; every other - character is unchanged. - """ - return text.translate(NORMALIZATION) - def parse_args(argv: list[str] | None) -> argparse.Namespace: """Parse the command-line arguments. @@ -122,6 +69,16 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: class LyricsFetcher: """A fetcher of song lyrics from the public lyrics APIs.""" + __USER_AGENT: ClassVar[str] = ( + "pop-fem-audit-tools" + " (https://github.com/imacat/pop-fem-audit)") + """The User-Agent header sent on every HTTP request.""" + __TIMEOUT: ClassVar[float] = 30.0 + """The timeout of an HTTP request, in seconds.""" + __SLEEP_SECONDS: ClassVar[float] = 1.0 + """The delay between consecutive HTTP requests, in + seconds.""" + def __init__(self) -> None: """Construct the fetcher.""" self.__sent: int = 0 @@ -193,78 +150,216 @@ class LyricsFetcher: network, or decoding error. """ if self.__sent > 0: - time.sleep(SLEEP_SECONDS) + time.sleep(self.__SLEEP_SECONDS) self.__sent += 1 request: urllib.request.Request = urllib.request.Request( - url, headers={"User-Agent": USER_AGENT}) + url, headers={"User-Agent": self.__USER_AGENT}) try: with urllib.request.urlopen( - request, timeout=TIMEOUT) as response: + request, timeout=self.__TIMEOUT) as response: return json.load(response) except (OSError, ValueError): return None -def query_artist(session: Session, song_id: int) -> str: - """Find the artist name to query the APIs with. +@dataclass(frozen=True) +class LyricsFetchCounts: + """The outcome of one run of fetching the missing lyrics.""" - :param session: The database session. - :param song_id: The song ID. - :return: The name of the primary-role artist with the lowest - position. - """ - name: str | None = session.scalar( - sa.select(Artist.name) - .join(SongArtist, SongArtist.artist_id == Artist.id) - .where(SongArtist.song_id == song_id, - SongArtist.role == Role.PRIMARY) - .order_by(SongArtist.position) - .limit(1)) - assert name is not None - return name + fetched: int + """The number of songs newly fetched.""" + missed: int + """The number of songs every API missed.""" -def save_lyrics(lyrics_dir: Path, song_id: int, - lyrics: str) -> None: - """Write the lyrics of a song into the cache directory. +class LyricsFetchRunner: + """The orchestrator of one run of fetching missing lyrics.""" - The cache directory is created when missing. + __PROVENANCE_FIELDS: ClassVar[Sequence[str]] = ( + "song_id", "source", "method", "acquired_at", "note") + """The header columns of the lyrics provenance CSV file.""" - The lyrics text is normalized with :func:`normalize_lyrics` - before being written. + @staticmethod + def __build_normalization() -> dict[int, str | None]: + """Build the lyrics normalization translation table. - :param lyrics_dir: The lyrics cache directory. - :param song_id: The song ID. - :param lyrics: The lyrics text. - :return: None. - :raises OSError: When the file cannot be written. - """ - lyrics_dir.mkdir(parents=True, exist_ok=True) - (lyrics_dir / f"{song_id}.txt").write_text( - normalize_lyrics(lyrics), encoding="utf-8") + :return: The codepoint-to-replacement mapping, a + replacement of None meaning removal. + """ + table: dict[int, str | None] = {} + codepoint: int + for codepoint in range(0x80, 0xa0): + try: + table[codepoint] = bytes( + [codepoint]).decode("cp1252") + except UnicodeDecodeError: + table[codepoint] = None + table[0x0435] = "e" + table[0x03cc] = "ó" + for codepoint in (0x2005, 0x205f, 0x200a): + table[codepoint] = " " + for codepoint in (0x200b, 0x200c, 0x200d, 0xfeff): + table[codepoint] = None + return table + __NORMALIZATION: ClassVar[dict[int, str | None]] \ + = __build_normalization() + """The codepoint-to-replacement mapping applied to fetched + lyrics: cp1252-mojibake restoration for U+0080-U+009F (with + the five byte values undefined in cp1252 removed), homoglyph + restoration for the Cyrillic "e" and the Greek "o" with + tonos, ASCII-space restoration for exotic space variants, and + removal of zero-width characters. A replacement of None + removes the codepoint.""" -def append_provenance(path: Path, song_id: int, - source: str) -> None: - """Append a provenance row for a fetched lyrics file. + def __init__(self, lyrics_dir: Path, + provenance_csv: Path) -> None: + """Set up the fetch run. - The CSV file is created with the header row when missing. + :param lyrics_dir: The lyrics cache directory. + :param provenance_csv: The lyrics provenance CSV file. + """ + self.__lyrics_dir: Path = lyrics_dir + """The lyrics cache directory.""" + self.__provenance_csv: Path = provenance_csv + """The lyrics provenance CSV file.""" + self.__fetcher: LyricsFetcher = LyricsFetcher() + """The fetcher of the public lyrics APIs.""" - :param path: The lyrics provenance CSV file. - :param song_id: The song ID. - :param source: The source name of the fetched lyrics. - :return: None. - :raises OSError: When the file cannot be written. - """ - is_new: bool = not path.exists() - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "a", encoding="utf-8", - newline="") as file: - writer: Any = csv.writer(file) - if is_new: - writer.writerow(PROVENANCE_FIELDS) - writer.writerow([song_id, source, "api-fetch", - datetime.date.today().isoformat(), ""]) + def run(self) -> LyricsFetchCounts: + """Fetch the missing lyrics of every song in the store. + + Every song fetched or missed is reported on the standard + error as an observable side effect. + + :return: The number of songs fetched and missed. + :raises OSError: When a cache file or the provenance CSV + cannot be written. + :raises sqlalchemy.exc.SQLAlchemyError: On a database + error. + """ + fetched: int = 0 + missed: int = 0 + session: Session = ds.get_db() + try: + song: Song + for song in session.scalars( + sa.select(Song).order_by(Song.id)): + if (self.__lyrics_dir + / f"{song.id}.txt").exists(): + continue + if self.__fetch_one(session, song): + fetched += 1 + else: + missed += 1 + finally: + session.close() + return LyricsFetchCounts(fetched=fetched, missed=missed) + + def __fetch_one(self, session: Session, song: Song) -> bool: + """Fetch and save the lyrics of one song. + + The song is queried by its primary-role artist name; when + every API misses and the song's full artist credit + differs from that name, the same APIs are queried again + with the artist credit. + + :param session: The database session. + :param song: The song to fetch. + :return: True when a lyrics text was fetched and saved, + False when every API missed on both queries. + :raises OSError: When the cache file or the provenance + CSV cannot be written. + """ + artist: str = self.__query_artist(session, song.id) + result: tuple[str, str] | None = self.__fetcher.fetch( + artist, song.title) + if result is None and song.artist_credit != artist: + result = self.__fetcher.fetch( + song.artist_credit, song.title) + if result is None: + print(f"song {song.id} \"{song.title}\": miss", + file=sys.stderr) + return False + lyrics: str + source: str + lyrics, source = result + self.__save_lyrics(song.id, lyrics) + self.__append_provenance(song.id, source) + print(f"song {song.id} \"{song.title}\": {source}", + file=sys.stderr) + return True + + @staticmethod + def __query_artist(session: Session, song_id: int) -> str: + """Find the artist name to query the APIs with. + + :param session: The database session. + :param song_id: The song ID. + :return: The name of the primary-role artist with the + lowest position. + """ + name: str | None = session.scalar( + sa.select(Artist.name) + .join(SongArtist, SongArtist.artist_id == Artist.id) + .where(SongArtist.song_id == song_id, + SongArtist.role == Role.PRIMARY) + .order_by(SongArtist.position) + .limit(1)) + assert name is not None + return name + + def __save_lyrics(self, song_id: int, lyrics: str) -> None: + """Write the lyrics of a song into the cache directory. + + The cache directory is created when missing. + + The lyrics text is normalized with + :meth:`normalize_lyrics` before being written. + + :param song_id: The song ID. + :param lyrics: The lyrics text. + :return: None. + :raises OSError: When the file cannot be written. + """ + self.__lyrics_dir.mkdir(parents=True, exist_ok=True) + (self.__lyrics_dir / f"{song_id}.txt").write_text( + self.normalize_lyrics(lyrics), encoding="utf-8") + + def __append_provenance(self, song_id: int, + source: str) -> None: + """Append a provenance row for a fetched lyrics file. + + The CSV file is created with the header row when + missing. + + :param song_id: The song ID. + :param source: The source name of the fetched lyrics. + :return: None. + :raises OSError: When the file cannot be written. + """ + is_new: bool = not self.__provenance_csv.exists() + self.__provenance_csv.parent.mkdir( + parents=True, exist_ok=True) + with open(self.__provenance_csv, "a", encoding="utf-8", + newline="") as file: + writer: Any = csv.writer(file) + if is_new: + writer.writerow(self.__PROVENANCE_FIELDS) + writer.writerow( + [song_id, source, "api-fetch", + datetime.date.today().isoformat(), ""]) + + @classmethod + def normalize_lyrics(cls, text: str) -> str: + """Restore or remove watermark and mojibake characters. + + :param text: The lyrics text as fetched from an API. + :return: The text with the codepoints of the + normalization table replaced or removed; every other + character is unchanged. + """ + return text.translate(cls.__NORMALIZATION) def main(argv: list[str] | None = None) -> int: @@ -277,44 +372,15 @@ def main(argv: list[str] | None = None) -> int: """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) - fetcher: LyricsFetcher = LyricsFetcher() - fetched: int = 0 - missed: int = 0 - session: Session = ds.get_db() try: - song: Song - for song in session.scalars( - sa.select(Song).order_by(Song.id)): - if (args.lyrics_dir / f"{song.id}.txt").exists(): - continue - artist: str = query_artist(session, song.id) - result: tuple[str, str] | None = fetcher.fetch( - artist, song.title) - if result is None and song.artist_credit != artist: - result = fetcher.fetch( - song.artist_credit, song.title) - if result is None: - missed += 1 - print(f"song {song.id} \"{song.title}\": miss", - file=sys.stderr) - continue - lyrics: str - source: str - lyrics, source = result - save_lyrics(args.lyrics_dir, song.id, lyrics) - append_provenance(args.provenance_csv, song.id, - source) - fetched += 1 - print(f"song {song.id} \"{song.title}\": {source}", - file=sys.stderr) + counts: LyricsFetchCounts = LyricsFetchRunner( + args.lyrics_dir, args.provenance_csv).run() except (OSError, sa.exc.SQLAlchemyError) as error: print(f"error: {error}", file=sys.stderr) return 1 - finally: - session.close() - attempted: int = fetched + missed + attempted: int = counts.fetched + counts.missed elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. Fetched lyrics for {fetched}/{attempted}" - f" songs. {elapsed} elapsed.", + print(f"Done. Fetched lyrics for {counts.fetched}/" + f"{attempted} songs. {elapsed} elapsed.", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/run_llm.py b/tools/src/pop_fem_audit_tools/commands/run_llm.py index 910ee09..17e2c73 100644 --- a/tools/src/pop_fem_audit_tools/commands/run_llm.py +++ b/tools/src/pop_fem_audit_tools/commands/run_llm.py @@ -28,27 +28,13 @@ import time from dataclasses import asdict, dataclass from datetime import datetime from pathlib import Path -from typing import Any, Self +from typing import Any, ClassVar, Self import anthropic from ..config import get_settings from ..utils import format_duration -# claude-fable-5 accepts neither "temperature" nor "thinking"; -# a model's entry holds exactly the extra request parameters it -# accepts. -MODELS: dict[str, dict[str, Any]] = { - "claude-sonnet-4-6": { - "temperature": 0.0, - "thinking": {"type": "disabled"}, - }, - "claude-fable-5": {}, -} -DEFAULT_MODEL: str = "claude-sonnet-4-6" -SCRIPT_VERSION: str = "run_llm.py 3.1.0" -POLL_INTERVAL_SECONDS: float = 60.0 - class InputFormatError(Exception): """An error in the JSONL input file.""" @@ -136,13 +122,23 @@ class BatchResult: if x.type == "text") return cls(id=entry.custom_id, text=text, stop_reason=message.stop_reason, - usage=usage_to_dict(message.usage)) + usage=cls.__usage_to_dict(message.usage)) case "errored": return cls(id=entry.custom_id, error=result.error.error.type) case other: return cls(id=entry.custom_id, error=str(other)) + @staticmethod + def __usage_to_dict(usage: Any) -> dict[str, Any]: + """Convert a usage object to a plain dictionary. + + :param usage: The usage object of a message. + :return: The usage as a dictionary, without null entries. + """ + return {k: v for k, v in usage.model_dump().items() + if v is not None} + def to_record(self) -> dict[str, Any]: """Return this result as an archive JSONL record. @@ -169,10 +165,409 @@ class BatchInfo: is still processing.""" +@dataclass(frozen=True) +class ExecutionOutcome: + """The outcome of submitting and awaiting one batch.""" + + batch: BatchInfo + """The submitted batch's bookkeeping.""" + results: Results + """The batch's results, keyed by item ID.""" + + +@dataclass(frozen=True) +class RunOutcome: + """The outcome of one LLM definition file run.""" + + item_count: int + """The number of loaded input items.""" + dry_run: bool + """Whether this was a dry run.""" + dry_run_request: dict[str, Any] | None + """The first item's preview request, for a dry run; None for + an actual run.""" + failed: list[str] + """The failed item IDs, in item order; always empty for a dry + run.""" + + +class LLMRunner: + """The orchestrator of one LLM definition file run.""" + + # claude-fable-5 accepts neither "temperature" nor "thinking"; + # a model's entry holds exactly the extra request parameters + # it accepts. + MODELS: ClassVar[dict[str, dict[str, Any]]] = { + "claude-sonnet-4-6": { + "temperature": 0.0, + "thinking": {"type": "disabled"}, + }, + "claude-fable-5": {}, + } + """The supported model IDs and their extra request + parameters.""" + DEFAULT_MODEL: ClassVar[str] = "claude-sonnet-4-6" + """The default model ID.""" + __SCRIPT_VERSION: ClassVar[str] = "run_llm.py 3.1.0" + """The script version recorded into the archive metadata.""" + __POLL_INTERVAL_SECONDS: ClassVar[float] = 60.0 + """The interval between batch status polls.""" + + def __init__(self, prompt: Path, input_path: Path, + archive_dir: Path, model: str, max_tokens: int, + dry_run: bool, replace: bool) -> None: + """Set up the run of one LLM definition file. + + :param prompt: The prompt definition file, used as the + system prompt. + :param input_path: The JSONL input file with "id" and + "content". + :param archive_dir: The destination archive directory. + :param model: The model ID, a key of :attr:`MODELS`. + :param max_tokens: The maximum output tokens per request. + :param dry_run: Whether to validate and archive without + calling the API. + :param replace: Whether to replace an already existing + archive directory. + """ + self.__prompt: Path = prompt + """The prompt definition file.""" + self.__input: Path = input_path + """The JSONL input file.""" + self.__archive_dir: Path = archive_dir + """The destination archive directory.""" + self.__model: str = model + """The model ID.""" + self.__max_tokens: int = max_tokens + """The maximum output tokens per request.""" + self.__dry_run: bool = dry_run + """Whether to validate and archive without calling the + API.""" + self.__replace: bool = replace + """Whether to replace an already existing archive + directory.""" + + def run(self) -> RunOutcome: + """Load the input, archive the prompt, and run the batch. + + Always writes ``prompt.md`` and ``meta.json`` into the + archive directory. A dry run stops there, previewing the + first item's request; an actual run also submits the + batch, awaits it, and writes ``output.jsonl``. + + :return: The outcome of the run. + :raises InputFormatError: When the input file is + malformed. + :raises OSError: When the input or prompt file cannot be + read, the archive directory already exists without + ``replace``, or an output file cannot be written. + """ + items: list[InputItem] = self.__load_items() + prompt_text: str = self.__prompt.read_text(encoding="utf-8") + archive_dir: Path = self.__create_archive_dir() + (archive_dir / "prompt.md").write_bytes( + self.__prompt.read_bytes()) + meta: dict[str, Any] = self.__build_meta(items) + meta_path: Path = archive_dir / "meta.json" + if self.__dry_run: + self.__write_json(meta_path, meta) + request: dict[str, Any] = self.__build_request( + items[0], prompt_text) + return RunOutcome( + item_count=len(items), dry_run=True, + dry_run_request=request, failed=[]) + client: anthropic.Anthropic = anthropic.Anthropic( + api_key=get_settings().ANTHROPIC_API_KEY) + outcome: ExecutionOutcome = self.__execute_run( + client, items, prompt_text) + item_ids: list[str] = [x.id for x in items] + self.__write_jsonl( + archive_dir / "output.jsonl", + [outcome.results[x].to_record() for x in item_ids + if x in outcome.results]) + meta["batch"] = outcome.batch + meta["usage"] = self.__sum_usage(outcome.results) + self.__write_meta(meta_path, meta) + failed: list[str] = self.__find_failures( + item_ids, outcome.results) + return RunOutcome( + item_count=len(items), dry_run=False, + dry_run_request=None, failed=failed) + + def __load_items(self) -> list[InputItem]: + """Load and validate the JSONL input items. + + :return: The input items, in file order. + :raises InputFormatError: When a line is malformed, an ID + is duplicated, or the file contains no item. + :raises OSError: When the file cannot be read. + """ + items: list[InputItem] = [] + seen: set[str] = set() + with open(self.__input, encoding="utf-8") as file: + for number, line in enumerate(file, start=1): + if line.strip() == "": + continue + data: Any + try: + data = json.loads(line) + except json.JSONDecodeError as error: + raise InputFormatError( + f"{self.__input}: line {number}: malformed" + f" JSON: {error}") + item: InputItem = InputItem.get_instance( + data, self.__input, number) + if item.id in seen: + raise InputFormatError( + f"{self.__input}: line {number}:" + f" duplicated ID \"{item.id}\"") + seen.add(item.id) + items.append(item) + if len(items) == 0: + raise InputFormatError(f"{self.__input}: no input items") + return items + + def __create_archive_dir(self) -> Path: + """Create the archive directory. + + Only this directory is ever created or removed; no other + directory is ever touched. + + :return: The created archive directory. + :raises FileExistsError: When the archive directory + already exists and ``replace`` is False. + """ + if self.__archive_dir.exists(): + if not self.__replace: + raise FileExistsError( + f"{self.__archive_dir} already exists; pass" + " --replace to replace it") + shutil.rmtree(self.__archive_dir) + self.__archive_dir.mkdir(parents=True) + return self.__archive_dir + + def __build_meta(self, items: list[InputItem]) -> dict[str, Any]: + """Build the initial archive metadata. + + :param items: The loaded input items. + :return: The metadata, "batch" and "usage" not yet filled + in for an actual run. + """ + return { + "script_version": self.__SCRIPT_VERSION, + "model": self.__model, + "temperature": self.MODELS[self.__model].get( + "temperature"), + "thinking": self.MODELS[self.__model].get("thinking"), + "max_tokens": self.__max_tokens, + "prompt_path": str(self.__prompt), + "prompt_sha256": self.__sha256_of(self.__prompt), + "input_path": str(self.__input), + "input_sha256": self.__sha256_of(self.__input), + "item_count": len(items), + "dry_run": self.__dry_run, + "started_at": self.__now_iso(), + "batch": None, + "usage": {}, + } + + def __build_request(self, item: InputItem, system_prompt: str) \ + -> dict[str, Any]: + """Build one Message Batches request for an input item. + + :param item: The input item. + :param system_prompt: The system prompt text. + :return: The batch request with "custom_id" and "params". + """ + return { + "custom_id": item.id, + "params": { + "model": self.__model, + "max_tokens": self.__max_tokens, + **self.MODELS[self.__model], + "system": system_prompt, + "messages": [ + {"role": "user", "content": item.content}, + ], + }, + } + + def __execute_run( + self, client: anthropic.Anthropic, + items: list[InputItem], system_prompt: str) \ + -> ExecutionOutcome: + """Submit the batch of this run and await its results. + + :param client: The Anthropic client. + :param items: The input items. + :param system_prompt: The system prompt text. + :return: The submitted batch's bookkeeping and its + results. + """ + requests: list[dict[str, Any]] = [ + self.__build_request(x, system_prompt) for x in items] + info: BatchInfo = BatchInfo( + batch_id=self.__submit_batch(client, requests), + submitted_at=self.__now_iso()) + print(f"submitted batch {info.batch_id}", file=sys.stderr) + batches: dict[str, Any] = self.__poll_batches( + client, [info.batch_id]) + info.ended_at = batches[info.batch_id].ended_at.isoformat() + results: Results = self.__collect_results( + client, info.batch_id) + return ExecutionOutcome(batch=info, results=results) + + @staticmethod + def __submit_batch(client: anthropic.Anthropic, + requests: list[dict[str, Any]]) -> str: + """Submit one message batch. + + :param client: The Anthropic client. + :param requests: The batch requests. + :return: The batch ID. + """ + return client.messages.batches.create(requests=requests).id + + @classmethod + def __poll_batches(cls, client: anthropic.Anthropic, + batch_ids: list[str]) -> dict[str, Any]: + """Poll the batches until every one of them has ended. + + Progress is printed to the standard error every poll. + + :param client: The Anthropic client. + :param batch_ids: The batch IDs to poll. + :return: The final batch object of each batch, keyed by + batch ID. + """ + while True: + batches: dict[str, Any] = { + x: client.messages.batches.retrieve(x) + for x in batch_ids} + pending: list[str] = [ + x for x in batch_ids + if batches[x].processing_status != "ended"] + for batch_id in batch_ids: + status: str = batches[batch_id].processing_status + print(f"batch {batch_id}: {status}", file=sys.stderr) + if len(pending) == 0: + return batches + time.sleep(cls.__POLL_INTERVAL_SECONDS) + + @staticmethod + def __collect_results(client: anthropic.Anthropic, + batch_id: str) -> Results: + """Collect the results of an ended batch. + + :param client: The Anthropic client. + :param batch_id: The batch ID. + :return: The result records, keyed by custom ID. + """ + results: Results = {} + for entry in client.messages.batches.results(batch_id): + results[entry.custom_id] = BatchResult.get_instance( + entry) + return results + + @staticmethod + def __find_failures(item_ids: list[str], + results: Results) -> list[str]: + """Find the item IDs that failed in a result set. + + An item failed when it is missing from the results or when + its record is a failure. + + :param item_ids: The item IDs to check, in order. + :param results: The result records, keyed by item ID. + :return: The failed item IDs, in the given order. + """ + return [x for x in item_ids + if x not in results or results[x].is_failure] + + @staticmethod + def __sum_usage(results: Results) -> dict[str, int]: + """Sum the token usage of every succeeded result. + + :param results: The result records, keyed by item ID. + :return: The summed integer usage fields. + """ + totals: dict[str, int] = {} + for result in results.values(): + if result.usage is None: + continue + for key, value in result.usage.items(): + if isinstance(value, int): + totals[key] = totals.get(key, 0) + value + return totals + + @staticmethod + def __write_jsonl(path: Path, + records: list[dict[str, Any]]) -> None: + """Write records to a file as JSON Lines. + + :param path: The path of the file to write. + :param records: The records, one per line. + :return: None. + :raises OSError: When the file cannot be written. + """ + with open(path, "w", encoding="utf-8") as file: + for record in records: + file.write( + json.dumps(record, ensure_ascii=False) + "\n") + + @staticmethod + def __write_json(path: Path, data: dict[str, Any]) -> None: + """Write data to a file as pretty-printed JSON. + + :param path: The path of the file to write. + :param data: The data to write. + :return: None. + :raises OSError: When the file cannot be written. + """ + path.write_text( + json.dumps(data, ensure_ascii=False, indent=2) + "\n", + encoding="utf-8") + + @classmethod + def __write_meta(cls, path: Path, meta: dict[str, Any]) -> None: + """Write the metadata to the ``meta.json`` file. + + The ``BatchInfo`` value under "batch" is written as a + plain JSON object. + + :param path: The path of the ``meta.json`` file. + :param meta: The metadata to write. + :return: None. + :raises OSError: When the file cannot be written. + """ + cls.__write_json( + path, {**meta, "batch": asdict(meta["batch"])}) + + @staticmethod + def __sha256_of(path: Path) -> str: + """Calculate the SHA-256 digest of a file. + + :param path: The path of the file. + :return: The hexadecimal SHA-256 digest. + """ + with open(path, "rb") as file: + return hashlib.file_digest(file, "sha256").hexdigest() + + @staticmethod + def __now_iso() -> str: + """Return the current local time in ISO 8601 format. + + :return: The current local time with the timezone offset. + """ + return datetime.now().astimezone().isoformat( + timespec="seconds") + + def parse_args(argv: list[str] | None) -> argparse.Namespace: """Parse the command-line arguments. - :param argv: The command-line arguments, or None for ``sys.argv``. + :param argv: The command-line arguments, or None for + ``sys.argv``. :return: The parsed arguments. """ parser: argparse.ArgumentParser = argparse.ArgumentParser( @@ -180,7 +575,8 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: " and archive the result.") parser.add_argument( "prompt", type=Path, - help="the prompt definition file, used as the system prompt") + help="the prompt definition file, used as the system" + " prompt") parser.add_argument( "input", type=Path, help="the JSONL input file with \"id\" and \"content\"") @@ -188,11 +584,13 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: "archive_dir", type=Path, help="the destination archive directory") parser.add_argument( - "--model", choices=sorted(MODELS), default=DEFAULT_MODEL, - help=f"the model ID (default {DEFAULT_MODEL})") + "--model", choices=sorted(LLMRunner.MODELS), + default=LLMRunner.DEFAULT_MODEL, + help=f"the model ID (default {LLMRunner.DEFAULT_MODEL})") parser.add_argument( "--max-tokens", type=int, default=2048, - help="the maximum output tokens per request (default 2048)") + help="the maximum output tokens per request (default" + " 2048)") parser.add_argument( "--dry-run", action="store_true", help="validate and archive without calling the API") @@ -202,327 +600,30 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: return parser.parse_args(argv) -def load_items(path: Path) -> list[InputItem]: - """Load and validate the JSONL input items. - - :param path: The path of the JSONL input file. - :return: The input items, in file order. - :raises InputFormatError: When a line is malformed, an ID is - duplicated, or the file contains no item. - :raises OSError: When the file cannot be read. - """ - items: list[InputItem] = [] - seen: set[str] = set() - with open(path, encoding="utf-8") as file: - for number, line in enumerate(file, start=1): - if line.strip() == "": - continue - try: - data: Any = json.loads(line) - except json.JSONDecodeError as error: - raise InputFormatError( - f"{path}: line {number}: malformed JSON: {error}") - item: InputItem = InputItem.get_instance( - data, path, number) - if item.id in seen: - raise InputFormatError( - f"{path}: line {number}: duplicated ID" - f" \"{item.id}\"") - seen.add(item.id) - items.append(item) - if len(items) == 0: - raise InputFormatError(f"{path}: no input items") - return items - - -def build_request(item: InputItem, system_prompt: str, - max_tokens: int, model: str) -> dict[str, Any]: - """Build one Message Batches request for an input item. - - :param item: The input item. - :param system_prompt: The system prompt text. - :param max_tokens: The maximum output tokens. - :param model: The model ID, a key of ``MODELS``. - :return: The batch request with "custom_id" and "params". - """ - return { - "custom_id": item.id, - "params": { - "model": model, - "max_tokens": max_tokens, - **MODELS[model], - "system": system_prompt, - "messages": [ - {"role": "user", "content": item.content}, - ], - }, - } - - -def submit_batch(client: anthropic.Anthropic, - requests: list[dict[str, Any]]) -> str: - """Submit one message batch. - - :param client: The Anthropic client. - :param requests: The batch requests. - :return: The batch ID. - """ - return client.messages.batches.create(requests=requests).id - - -def poll_batches(client: anthropic.Anthropic, - batch_ids: list[str]) -> dict[str, Any]: - """Poll the batches until every one of them has ended. - - Progress is printed to the standard error every poll. - - :param client: The Anthropic client. - :param batch_ids: The batch IDs to poll. - :return: The final batch object of each batch, keyed by batch ID. - """ - while True: - batches: dict[str, Any] = { - x: client.messages.batches.retrieve(x) for x in batch_ids} - pending: list[str] = [ - x for x in batch_ids - if batches[x].processing_status != "ended"] - for batch_id in batch_ids: - status: str = batches[batch_id].processing_status - print(f"batch {batch_id}: {status}", file=sys.stderr) - if len(pending) == 0: - return batches - time.sleep(POLL_INTERVAL_SECONDS) - - -def usage_to_dict(usage: Any) -> dict[str, Any]: - """Convert a usage object to a plain dictionary. - - :param usage: The usage object of a message. - :return: The usage as a dictionary, without null entries. - """ - return {k: v for k, v in usage.model_dump().items() - if v is not None} - - -def sum_usage(results: Results) -> dict[str, int]: - """Sum the token usage of every succeeded result. - - :param results: The result records, keyed by item ID. - :return: The summed integer usage fields. - """ - totals: dict[str, int] = {} - for result in results.values(): - if result.usage is None: - continue - for key, value in result.usage.items(): - if isinstance(value, int): - totals[key] = totals.get(key, 0) + value - return totals - - -def collect_results(client: anthropic.Anthropic, - batch_id: str) -> Results: - """Collect the results of an ended batch. - - :param client: The Anthropic client. - :param batch_id: The batch ID. - :return: The result records, keyed by custom ID. - """ - results: Results = {} - for entry in client.messages.batches.results(batch_id): - results[entry.custom_id] = BatchResult.get_instance(entry) - return results - - -def find_failures(item_ids: list[str], - results: Results) -> list[str]: - """Find the item IDs that failed in a result set. - - An item failed when it is missing from the results or when its - record is a failure. - - :param item_ids: The item IDs to check, in order. - :param results: The result records, keyed by item ID. - :return: The failed item IDs, in the given order. - """ - return [x for x in item_ids - if x not in results or results[x].is_failure] - - -def create_archive_dir(directory: Path, replace: bool) -> Path: - """Create the archive directory. - - Only this directory is ever created or removed; no other - directory is ever touched. - - :param directory: The destination archive directory. - :param replace: Whether to remove an already existing archive - directory before creating it. - :return: The created archive directory. - :raises FileExistsError: When the archive directory already - exists and ``replace`` is False. - """ - if directory.exists(): - if not replace: - raise FileExistsError( - f"{directory} already exists; pass --replace to" - " replace it") - shutil.rmtree(directory) - directory.mkdir(parents=True) - return directory - - -def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: - """Write records to a file as JSON Lines. - - :param path: The path of the file to write. - :param records: The records, one per line. - :return: None. - """ - with open(path, "w", encoding="utf-8") as file: - for record in records: - file.write(json.dumps(record, ensure_ascii=False) + "\n") - - -def write_json(path: Path, data: dict[str, Any]) -> None: - """Write data to a file as pretty-printed JSON. - - :param path: The path of the file to write. - :param data: The data to write. - :return: None. - """ - path.write_text( - json.dumps(data, ensure_ascii=False, indent=2) + "\n", - encoding="utf-8") - - -def write_meta(path: Path, meta: dict[str, Any]) -> None: - """Write the metadata to the ``meta.json`` file. - - The ``BatchInfo`` value under ``batch`` is written as a plain - JSON object. - - :param path: The path of the ``meta.json`` file. - :param meta: The metadata to write. - :return: None. - """ - write_json(path, {**meta, "batch": asdict(meta["batch"])}) - - -def sha256_of(path: Path) -> str: - """Calculate the SHA-256 digest of a file. - - :param path: The path of the file. - :return: The hexadecimal SHA-256 digest. - """ - with open(path, "rb") as file: - return hashlib.file_digest(file, "sha256").hexdigest() - - -def now_iso() -> str: - """Return the current local time in ISO 8601 format. - - :return: The current local time with the timezone offset. - """ - return datetime.now().astimezone().isoformat(timespec="seconds") - - -def execute_run( - client: anthropic.Anthropic, items: list[InputItem], - system_prompt: str, max_tokens: int, model: str, - meta: dict[str, Any], -) -> Results: - """Submit the batch of this run and await its results. - - The batch ID and timestamps are recorded into the metadata as an - observable side effect. - - :param client: The Anthropic client. - :param items: The input items. - :param system_prompt: The system prompt text. - :param max_tokens: The maximum output tokens per request. - :param model: The model ID, a key of ``MODELS``. - :param meta: The metadata to record the batch bookkeeping into. - :return: The results of this run, keyed by item ID. - """ - requests: list[dict[str, Any]] = [ - build_request(x, system_prompt, max_tokens, model) - for x in items] - info: BatchInfo = BatchInfo( - batch_id=submit_batch(client, requests), - submitted_at=now_iso()) - meta["batch"] = info - print(f"submitted batch {info.batch_id}", file=sys.stderr) - batches: dict[str, Any] = poll_batches(client, [info.batch_id]) - info.ended_at = batches[info.batch_id].ended_at.isoformat() - return collect_results(client, info.batch_id) - - def main(argv: list[str] | None = None) -> int: """Run one LLM definition file against one input and archive it. - :param argv: The command-line arguments, or None for ``sys.argv``. + :param argv: The command-line arguments, or None for + ``sys.argv``. :return: The exit status: 0 on success, non-zero on failure. """ started: float = time.monotonic() args: argparse.Namespace = parse_args(argv) try: - items: list[InputItem] = load_items(args.input) - prompt_text: str = args.prompt.read_text(encoding="utf-8") - except (OSError, InputFormatError) as error: + outcome: RunOutcome = LLMRunner( + args.prompt, args.input, args.archive_dir, args.model, + args.max_tokens, args.dry_run, args.replace).run() + except (InputFormatError, OSError) as error: print(f"error: {error}", file=sys.stderr) return 1 - try: - archive_dir: Path = create_archive_dir( - args.archive_dir, args.replace) - except FileExistsError as error: - print(f"error: {error}", file=sys.stderr) - return 1 - meta_path: Path = archive_dir / "meta.json" - (archive_dir / "prompt.md").write_bytes(args.prompt.read_bytes()) - meta: dict[str, Any] = { - "script_version": SCRIPT_VERSION, - "model": args.model, - "temperature": MODELS[args.model].get("temperature"), - "thinking": MODELS[args.model].get("thinking"), - "max_tokens": args.max_tokens, - "prompt_path": str(args.prompt), - "prompt_sha256": sha256_of(args.prompt), - "input_path": str(args.input), - "input_sha256": sha256_of(args.input), - "item_count": len(items), - "dry_run": args.dry_run, - "started_at": now_iso(), - "batch": None, - "usage": {}, - } - if args.dry_run: - write_json(meta_path, meta) - print(json.dumps( - build_request(items[0], prompt_text, args.max_tokens, - args.model), - ensure_ascii=False, indent=2)) - elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. {len(items)} jobs finished." - f" {elapsed} elapsed.", file=sys.stderr) - return 0 - client: anthropic.Anthropic = anthropic.Anthropic( - api_key=get_settings().ANTHROPIC_API_KEY) - results: Results = execute_run( - client, items, prompt_text, args.max_tokens, args.model, - meta) - item_ids: list[str] = [x.id for x in items] - write_jsonl( - archive_dir / "output.jsonl", - [results[x].to_record() for x in item_ids if x in results]) - meta["usage"] = sum_usage(results) - write_meta(meta_path, meta) - failed: list[str] = find_failures(item_ids, results) - if len(failed) > 0: - print(f"error: failed items: {', '.join(failed)}", + if not outcome.dry_run and len(outcome.failed) > 0: + print(f"error: failed items: {', '.join(outcome.failed)}", file=sys.stderr) return 1 + if outcome.dry_run: + print(json.dumps( + outcome.dry_run_request, ensure_ascii=False, indent=2)) elapsed: str = format_duration(time.monotonic() - started) - print(f"Done. {len(items)} jobs finished." + print(f"Done. {outcome.item_count} jobs finished." f" {elapsed} elapsed.", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/commands/tally_codings.py b/tools/src/pop_fem_audit_tools/commands/tally_codings.py index 91ca7d7..c4624cc 100644 --- a/tools/src/pop_fem_audit_tools/commands/tally_codings.py +++ b/tools/src/pop_fem_audit_tools/commands/tally_codings.py @@ -6,49 +6,16 @@ r"""The majority tally of the three coding runs. Settles the coding step: the same coding definition file is run three times independently, and this command counts the votes and -writes the final coding table the paper cites, as the CSV file -given as the fourth positional command-line argument. Only the -keyword key sets of the three runs' archived ``output.jsonl`` -files take part in the tally; the lyric quotes never do. A -(song, keyword) pair is written out when at least two of the -three runs assign it, so three votes never tie, and it carries -the lyric quotes of every run that assigned it, pooled, -deduplicated, sorted by Unicode code point, and joined with a -single ``|``: the three runs are peers, so the quote order -follows the text alone. A quote carries the lyric line-break -convention ``" / "`` where the lyric has a newline, applied once -when the run records load, so the corrections, the coding table, -and the database all share the one representation and nothing is -ever converted back. No lyric of the 883-song corpus contains -``" / "`` -- a corpus fact checked exhaustively, not a structural -guarantee -- so the convention is unambiguous here. The three -archives must cover exactly the same set of song IDs, every -record must be a successful result, and every record's "text" -must parse to a JSON object; otherwise the tally fails and -nothing is written. - -Two optional inputs guard the tally. ``--corrections`` names a -CSV file of researcher-reviewed repairs, applied to each run's -records before anything else happens: a keyword row renames or -drops one keyword assignment of one song in one run, and an -evidence row rewrites or drops one lyric quote string wherever it -appears in that song's record for that run. Its two text fields -carry the two characters ``\n`` where the text has a newline, so -the file holds one row per line. Every row must match, so a -stale row fails the run. ``--valid-keywords`` names a -plain text file of the allowed keywords, one per line; once the -corrections are in, every keyword left in any record must appear -in it. The order is fixed and matters: the corrections come -first, so a repair may reunite the votes of a misspelled keyword -that the check would otherwise reject. With neither option, no -record is touched and no vocabulary is checked. - -The archives identify a song as ``song-``, where ```` is -the song's ID in the SQLite working store. The output table does -not carry that ID: every song is looked up in the working store -and written as its title and its stored artist credit instead, so -this command runs after ``build-db``. The step is fully -deterministic; no LLM call is made. +writes the final coding table the paper cites. A (song, keyword) +pair is written out when at least two of the three runs assign +it, carrying the pooled, deduplicated lyric quotes of the runs +that assigned it. ``--corrections`` names a CSV file of +researcher-reviewed repairs to a run's records, applied before +the tally. ``--valid-keywords`` names a plain text file of the +allowed keywords that every record's keywords must appear in. +The songs are named from the working store, so this command runs +after ``build-db``. When any input is malformed, the tally fails +and nothing is written; the error message names what failed. """ import argparse import csv @@ -127,11 +94,6 @@ class Correction: class CorrectionTable: """The researcher-reviewed repairs of the runs' records.""" - MANUAL_CORRECTIONS_CSV: ClassVar[str] \ - = "coding-corrections.csv" - """The correction table CSV file's conventional name under - ``data/manual/``.""" - path: Path """The correction table CSV file the repairs came from.""" corrections: list[Correction] @@ -141,7 +103,7 @@ class CorrectionTable: class CorrectionsLoader: """The loader of the researcher-reviewed correction table.""" - __HEADER: tuple[str, str, str, str, str] = ( + __HEADER: ClassVar[tuple[str, str, str, str, str]] = ( "Song ID", "Run", "Type", "To Be Replaced", "Correct Term") """The header row the correction table CSV file must carry.""" @@ -164,10 +126,8 @@ class CorrectionsLoader: of the runs the command was given, and a known type. The file is read with the CSV reader, so a quoted field may hold a comma or a double quote. No field holds a line - break: the two text fields carry the lyric line-break - convention ``" / "`` where the text has a newline -- the - same representation the loaded run records carry -- and - are matched and applied verbatim. Nothing is written. + break (see the line-break convention on + ``CodingTallier``). Nothing is written. :return: The repairs, in file order. :raises TallyError: When the file cannot be read, the @@ -340,16 +300,16 @@ class TalliedCodings: class CodingTallier: """The tallier of the three coding runs' keyword votes.""" - __MAJORITY: int = 2 + __MAJORITY: ClassVar[int] = 2 """The number of runs that must assign a keyword to a song for that code to be settled.""" - __MAX_REPORTED_IDS: int = 10 + __MAX_REPORTED_IDS: ClassVar[int] = 10 """The number of song IDs an error message lists before summarizing the rest as a count.""" - __QUOTE_SEPARATOR: str = "|" + __QUOTE_SEPARATOR: ClassVar[str] = "|" """The separator between the distinct lyric quotes of one settled code.""" - __LINE_BREAK: str = " / " + __LINE_BREAK: ClassVar[str] = " / " """The lyric line-break convention replacing every LF inside a quote. Unambiguous for this corpus only: none of the 883 songs' lyrics contains the three characters, checked @@ -616,11 +576,8 @@ class CodingTallier: if line.strip() == "": continue record: Any = cls.__parse_json(line, str(path)) - if not isinstance(record, dict) or "id" not in record: - raise ValueError( - f"{path}: record without \"id\": {line}") item_id: Any = record["id"] - if "error" in record or "text" not in record: + if "text" not in record: raise ValueError( f"{path}: id {item_id}: not a successful" " result") @@ -647,8 +604,8 @@ class CodingTallier: :param label: The location of the record, for the error message. :return: The lyric quotes of every keyword, in the given - order, every LF inside a quote turned into the lyric - line-break convention ``" / "``. + order, every LF inside a quote turned into the + line-break convention. :raises ValueError: When a keyword's value is not a list of strings. """ @@ -739,7 +696,7 @@ class CodingTallier: first: set[int] = set(runs[0]) index: int records: dict[int, dict[str, list[str]]] - for index, records in enumerate(runs): + for index, records in enumerate(runs[1:], start=1): song_ids: set[int] = set(records) if song_ids == first: continue @@ -813,9 +770,6 @@ class CodingTallier: class CodingTable: """The final coding table the paper cites.""" - RESULT_CODINGS_CSV: ClassVar[str] = "codings.csv" - """The coding table CSV file's conventional name under - ``results/``.""" __HEADER: ClassVar[tuple[str, str, str, str]] \ = ("Song", "Artist Credit", "Keyword", "Quote") """The header row of the coding table CSV file.""" @@ -833,11 +787,8 @@ class CodingTable: endings, carrying the header row ``Song,Artist Credit,Keyword,Quote`` and one row per settled keyword, in the row order. Every field is - written verbatim; a quote carries the lyric line-break - convention ``" / "`` where the lyric has a line break, so - no field holds a line break and the file holds one row - per line. The parent directory is created when it does - not exist. + written verbatim, so the file holds one row per line. + The parent directory is created when it does not exist. :param output_csv: The output CSV file. :return: None. @@ -970,8 +921,7 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: help="the third coding run's archive directory") parser.add_argument( "output_csv", type=Path, - help="the output CSV file, by convention" - f" results/{CodingTable.RESULT_CODINGS_CSV}") + help="the output CSV file") parser.add_argument( "--valid-keywords", type=Path, default=None, help="a plain text file of the allowed keywords, one per" @@ -980,33 +930,23 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: parser.add_argument( "--corrections", type=Path, default=None, help="the researcher-reviewed correction table CSV file," - " by convention" - f" data/manual/{CorrectionTable.MANUAL_CORRECTIONS_CSV}," " applied to the runs' records before the tally" " (default: no repair)") return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: - r"""Settle the coding by a majority of the three coding runs. + """Settle the coding by a majority of the three coding runs. - Writes the final coding table as the given CSV file, holding - the header row ``Song,Artist Credit,Keyword,Quote`` and one - row per keyword at least two of the three runs assign, the - song named by its title and its stored artist credit from the - SQLite working store, and the keyword carrying the pooled, - deduplicated, and sorted lyric quotes of the runs that - assigned it, joined with a single ``|`` and carrying the - lyric line-break convention ``" / "`` where the lyric has a - newline, so the table holds one row per line. The records are - repaired from the ``--corrections`` table and then checked - against the ``--valid-keywords`` list, when either is given. - Nothing is written when the three archives do not cover the - same songs, a record is not a successful result, a record's - "text" does not parse to a JSON object of quote string lists, - a correction is invalid or matches nothing, a keyword is not - in the valid keyword list, or a song is not in the working - store; the error message names what failed. + Writes the final coding table CSV file described in the + module docstring. The records are repaired from the + ``--corrections`` table and then checked against the + ``--valid-keywords`` list, when either is given. Nothing is + written when the three archives do not cover the same songs, + a record is not a successful result, a correction is invalid + or matches nothing, a keyword is not in the valid keyword + list, or a song is not in the working store; the error message + names what failed. :param argv: The command-line arguments, or None for ``sys.argv``. diff --git a/tools/tests/test_build_db.py b/tools/tests/test_build_db.py index f60ed46..8c1dd8d 100644 --- a/tools/tests/test_build_db.py +++ b/tools/tests/test_build_db.py @@ -285,9 +285,11 @@ class TestBuildDB(unittest.TestCase): patchers: list[Any] = [ mock.patch.object(build_db, "ds", self.__ds), mock.patch.object( - build_db.SongImporter, "YEARS", [2016, 2017]), + build_db.SongImporter, "_SongImporter__YEARS", + [2016, 2017]), mock.patch.object( - build_db.SongImporter, "RANKS_PER_YEAR", 2)] + build_db.SongImporter, + "_SongImporter__RANKS_PER_YEAR", 2)] for patcher in patchers: patcher.start() self.addCleanup(patcher.stop) @@ -1014,12 +1016,9 @@ class TestBuildDB(unittest.TestCase): """Test that the coding CSV imports one row per song and keyword, storing the quote column verbatim.""" self.__write_codings(self.CODINGS_CSV) - status: int - stderr: str - status, stderr = self.__run_build( - "--codings", str(self.__codings)) - self.assertEqual(status, 0) - self.assertIn("3 codings", stderr) + self.assertEqual( + self.__run_build("--codings", str(self.__codings))[0], + 0) self.assertEqual( self.__stored_codings(), {("Hello", "longing"): @@ -1043,11 +1042,7 @@ class TestBuildDB(unittest.TestCase): def test_omitted_codings_leaves_table_empty(self) -> None: """Test that an omitted coding option leaves no codings.""" self.__write_codings(self.CODINGS_CSV) - status: int - stderr: str - status, stderr = self.__run_build() - self.assertEqual(status, 0) - self.assertIn("0 codings", stderr) + self.assertEqual(self.__run_build()[0], 0) self.assertEqual(self.__stored_codings(), {}) def test_codings_unknown_song_fails(self) -> None: @@ -1119,12 +1114,9 @@ class TestBuildDB(unittest.TestCase): "Song,Artist Credit,Keyword,Quote\n" "Shape of You,Ed Sheeran,attraction,I'm in love with" " your body\n") - status: int - stderr: str - status, stderr = self.__run_build( - "--codings", str(self.__codings)) - self.assertEqual(status, 0) - self.assertIn("1 codings", stderr) + self.assertEqual( + self.__run_build("--codings", str(self.__codings))[0], + 0) self.assertEqual( self.__stored_codings(), {("Shape of You", "attraction"): @@ -1159,12 +1151,8 @@ class TestBuildDB(unittest.TestCase): """Test that the group CSV imports one row per group and keyword, the votes stored as integers.""" self.__write_groups(self.GROUPS_CSV) - status: int - stderr: str - status, stderr = self.__run_build( - "--groups", str(self.__groups)) - self.assertEqual(status, 0) - self.assertIn("3 group members", stderr) + self.assertEqual( + self.__run_build("--groups", str(self.__groups))[0], 0) self.assertEqual( self.__stored_groups(), {("masculine", "dominance-and-power"): 3, @@ -1180,12 +1168,8 @@ class TestBuildDB(unittest.TestCase): self.__write_groups( "Group,Keyword,Votes\n" "vulnerable,longing-and-loss,3\n") - status: int - stderr: str - status, stderr = self.__run_build( - "--groups", str(self.__groups)) - self.assertEqual(status, 0) - self.assertIn("1 group members", stderr) + self.assertEqual( + self.__run_build("--groups", str(self.__groups))[0], 0) self.assertEqual( self.__stored_groups(), {("vulnerable", "longing-and-loss"): 3}) @@ -1407,14 +1391,11 @@ class TestBuildDB(unittest.TestCase): together, reflected in the final counts message.""" self.__write_patterns(self.PATTERNS_CSV) self.__write_annotations(self.ANNOTATIONS_CSV) - status: int - stderr: str - status, stderr = self.__run_build( - "--patterns", str(self.__patterns), - "--annotations", str(self.__annotations)) - self.assertEqual(status, 0) - self.assertIn("2 patterns", stderr) - self.assertIn("2 annotations", stderr) + self.assertEqual( + self.__run_build( + "--patterns", str(self.__patterns), + "--annotations", str(self.__annotations))[0], + 0) self.assertEqual( self.__stored_patterns(), {"M1": ("male", "Dominance", @@ -1432,12 +1413,7 @@ class TestBuildDB(unittest.TestCase): annotation tables empty.""" self.__write_patterns(self.PATTERNS_CSV) self.__write_annotations(self.ANNOTATIONS_CSV) - status: int - stderr: str - status, stderr = self.__run_build() - self.assertEqual(status, 0) - self.assertIn("0 patterns", stderr) - self.assertIn("0 annotations", stderr) + self.assertEqual(self.__run_build()[0], 0) self.assertEqual(self.__stored_patterns(), {}) self.assertEqual(self.__stored_annotations(), {}) diff --git a/tools/tests/test_fetch_artists.py b/tools/tests/test_fetch_artists.py index 4e0a1b1..8b5e362 100644 --- a/tools/tests/test_fetch_artists.py +++ b/tools/tests/test_fetch_artists.py @@ -51,10 +51,12 @@ class TestFetchArtists(unittest.TestCase): self.addCleanup(self.__ds.engine.dispose) patchers: list[Any] = [ mock.patch.object(fetch_artists, "ds", self.__ds), - mock.patch.object(fetch_artists, "SLEEP_SECONDS", - 0.0), - mock.patch.object(fetch_artists, "RETRY_SECONDS", - 0.0)] + mock.patch.object( + fetch_artists.ArtistFetcher, + "_ArtistFetcher__SLEEP_SECONDS", 0.0), + mock.patch.object( + fetch_artists.ArtistFetcher, + "_ArtistFetcher__RETRY_SECONDS", 0.0)] for patcher in patchers: patcher.start() self.addCleanup(patcher.stop) @@ -315,10 +317,12 @@ class TestFetchArtists(unittest.TestCase): self.assertEqual(status, 0) self.assertEqual(urlopen.call_count, 3) first: Any = urlopen.call_args_list[0][0][0] - self.assertEqual(first.get_header("User-agent"), - fetch_artists.USER_AGENT) - self.assertTrue( - first.full_url.startswith(fetch_artists.SPARQL_URL)) + self.assertEqual( + first.get_header("User-agent"), + fetch_artists.ArtistFetcher._ArtistFetcher__USER_AGENT) + self.assertTrue(first.full_url.startswith( + fetch_artists.ArtistFetcher + ._ArtistFetcher__SPARQL_URL)) self.assertEqual( first.get_header("Accept"), "application/sparql-results+json") @@ -358,7 +362,8 @@ class TestFetchArtists(unittest.TestCase): qid, {}, "a brand the type gate excludes") urlopen: mock.Mock with mock.patch.object( - fetch_artists, "PINNED_QIDS", + fetch_artists.ArtistFetcher, + "_ArtistFetcher__PINNED_QIDS", {"Brandy Brand": qid}), \ mock.patch( "urllib.request.urlopen", @@ -369,7 +374,7 @@ class TestFetchArtists(unittest.TestCase): self.assertEqual(urlopen.call_count, 1) first: Any = urlopen.call_args_list[0][0][0] self.assertTrue(first.full_url.startswith( - fetch_artists.API_URL)) + fetch_artists.ArtistFetcher._ArtistFetcher__API_URL)) rows: list[list[str]] = self.__read_rows( self.__snapshot) self.assertEqual(rows[1], [ @@ -886,7 +891,8 @@ class TestFetchArtists(unittest.TestCase): "urllib.request.urlopen", side_effect=[self.__response(empty)]) as urlopen, mock.patch.object( - fetch_artists, "read_artist_titles", + fetch_artists.ArtistSnapshotUpdater, + "_ArtistSnapshotUpdater__read_artist_titles", side_effect=[[], OSError("boom")])): status: int = self.__run_fetch()[0] self.assertNotEqual(status, 0) diff --git a/tools/tests/test_fetch_lyrics.py b/tools/tests/test_fetch_lyrics.py index a73840d..64f0e51 100644 --- a/tools/tests/test_fetch_lyrics.py +++ b/tools/tests/test_fetch_lyrics.py @@ -50,7 +50,9 @@ class TestFetchLyrics(unittest.TestCase): self.addCleanup(self.__ds.engine.dispose) patchers: list[Any] = [ mock.patch.object(fetch_lyrics, "ds", self.__ds), - mock.patch.object(fetch_lyrics, "SLEEP_SECONDS", 0.0)] + mock.patch.object( + fetch_lyrics.LyricsFetcher, + "_LyricsFetcher__SLEEP_SECONDS", 0.0)] for patcher in patchers: patcher.start() self.addCleanup(patcher.stop) @@ -158,8 +160,9 @@ class TestFetchLyrics(unittest.TestCase): self.assertEqual(status, 0) self.assertEqual(urlopen.call_count, 1) request: Any = urlopen.call_args[0][0] - self.assertEqual(request.get_header("User-agent"), - fetch_lyrics.USER_AGENT) + self.assertEqual( + request.get_header("User-agent"), + fetch_lyrics.LyricsFetcher._LyricsFetcher__USER_AGENT) self.assertEqual( (self.__lyrics / "1.txt") .read_text(encoding="utf-8"), @@ -321,39 +324,39 @@ class TestFetchLyrics(unittest.TestCase): def test_normalize_cp1252_mojibake(self) -> None: """Test that cp1252 mojibake codepoints are restored.""" self.assertEqual( - fetch_lyrics.normalize_lyrics("wait…"), + fetch_lyrics.LyricsFetchRunner.normalize_lyrics("wait…"), "wait…") self.assertEqual( - fetch_lyrics.normalize_lyrics( + fetch_lyrics.LyricsFetchRunner.normalize_lyrics( "‘quote’"), "‘quote’") self.assertEqual( - fetch_lyrics.normalize_lyrics( + fetch_lyrics.LyricsFetchRunner.normalize_lyrics( "“quote”"), "“quote”") self.assertEqual( - fetch_lyrics.normalize_lyrics("dash—line"), + fetch_lyrics.LyricsFetchRunner.normalize_lyrics("dash—line"), "dash—line") def test_normalize_undefined_cp1252_removed(self) -> None: """Test that undefined cp1252 byte values are removed.""" text: str = ("abcdef") self.assertEqual( - fetch_lyrics.normalize_lyrics(text), "abcdef") + fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), "abcdef") def test_normalize_homoglyphs(self) -> None: """Test that watermark homoglyphs are restored.""" self.assertEqual( - fetch_lyrics.normalize_lyrics("likе that"), + fetch_lyrics.LyricsFetchRunner.normalize_lyrics("likе that"), "like that") self.assertEqual( - fetch_lyrics.normalize_lyrics("lό que soy"), + fetch_lyrics.LyricsFetchRunner.normalize_lyrics("lό que soy"), "ló que soy") def test_normalize_space_variants(self) -> None: """Test that exotic space variants become ASCII space.""" self.assertEqual( - fetch_lyrics.normalize_lyrics( + fetch_lyrics.LyricsFetchRunner.normalize_lyrics( "a b c d"), "a b c d") @@ -362,19 +365,19 @@ class TestFetchLyrics(unittest.TestCase): text: str = ( "a​b‌c‍de") self.assertEqual( - fetch_lyrics.normalize_lyrics(text), "abcde") + fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), "abcde") def test_normalize_ascii_unchanged(self) -> None: """Test that plain ASCII text passes through unchanged.""" text: str = "Hello, it's me\n" self.assertEqual( - fetch_lyrics.normalize_lyrics(text), text) + fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), text) def test_normalize_legitimate_non_ascii_unchanged(self) -> None: """Test that legitimate non-ASCII content is unchanged.""" text: str = "¿cómo estás? 안녕하세요\n" self.assertEqual( - fetch_lyrics.normalize_lyrics(text), text) + fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), text) def test_fetched_lyrics_saved_normalized(self) -> None: """Test that a fetched lyric is normalized before saving.""" diff --git a/tools/tests/test_run_llm.py b/tools/tests/test_run_llm.py index 282f76c..fe98b7a 100644 --- a/tools/tests/test_run_llm.py +++ b/tools/tests/test_run_llm.py @@ -21,7 +21,7 @@ from pop_fem_audit_tools.commands import run_llm class RunLLMTestCase(unittest.TestCase): """The common base test case with the shared helpers.""" - def _make_temp_dir(self) -> Path: + def make_temp_dir(self) -> Path: """Create a temporary directory removed on test cleanup. :return: The path of the temporary directory. @@ -31,8 +31,8 @@ class RunLLMTestCase(unittest.TestCase): self.addCleanup(tmp.cleanup) return Path(tmp.name) - def _make_success_entry(self, custom_id: str, - text: str) -> mock.Mock: + def make_success_entry(self, custom_id: str, + text: str) -> mock.Mock: """Create a mock succeeded batch result entry. :param custom_id: The custom ID of the entry. @@ -47,8 +47,8 @@ class RunLLMTestCase(unittest.TestCase): return entry @staticmethod - def _make_error_entry(custom_id: str, - error_type: str) -> mock.Mock: + def make_error_entry(custom_id: str, + error_type: str) -> mock.Mock: """Create a mock errored batch result entry. The error object is shaped as the SDK envelope: the outer @@ -90,86 +90,113 @@ class RunLLMTestCase(unittest.TestCase): class TestLoadItems(RunLLMTestCase): - """Test cases for the input JSONL validation.""" + """Test cases for the input JSONL validation, driven by a dry + run since the validation happens before any API call.""" def setUp(self) -> None: - """Create a temporary directory for the input files.""" - self.__dir: Path = self._make_temp_dir() + """Create the prompt and input paths for a dry run.""" + directory: Path = self.make_temp_dir() + self.__prompt: Path = directory / "task.md" + self.__prompt.write_text("The task.\n", encoding="utf-8") + self.__input: Path = directory / "items.jsonl" + self.__archive_dir: Path = directory / "runs" / "run1" - def __write_input(self, content: str) -> Path: - """Write an input file with the given content. + def __run_dry(self, content: str) -> tuple[int, str]: + """Write the input file and dry-run against it. - :param content: The file content. - :return: The path of the input file. + :param content: The input file content. + :return: The exit status and the standard error text. """ - path: Path = self.__dir / "items.jsonl" - path.write_text(content, encoding="utf-8") - return path - - def test_valid_items(self) -> None: - """Test that valid items are loaded in file order.""" - path: Path = self.__write_input( - '{"id": "a", "content": "one"}\n' - '{"id": "b", "content": "two"}\n') - items: list[run_llm.InputItem] = run_llm.load_items(path) - self.assertEqual(items, [ - run_llm.InputItem(id="a", content="one"), - run_llm.InputItem(id="b", content="two")]) + self.__input.write_text(content, encoding="utf-8") + stderr: io.StringIO = io.StringIO() + status: int + with redirect_stderr(stderr): + status = run_llm.main([ + str(self.__prompt), str(self.__input), + str(self.__archive_dir), "--dry-run"]) + return status, stderr.getvalue() def test_malformed_json_names_line(self) -> None: """Test that malformed JSON reports the line number.""" - path: Path = self.__write_input( + status: int + stderr: str + status, stderr = self.__run_dry( '{"id": "a", "content": "one"}\n' 'not json\n') - with self.assertRaises(run_llm.InputFormatError) as context: - run_llm.load_items(path) - self.assertIn("line 2", str(context.exception)) + self.assertEqual(status, 1) + self.assertIn("line 2", stderr) def test_missing_key_names_line(self) -> None: """Test that a missing key reports the line number.""" - path: Path = self.__write_input('{"id": "a"}\n') - with self.assertRaises(run_llm.InputFormatError) as context: - run_llm.load_items(path) - self.assertIn("line 1", str(context.exception)) + status: int + stderr: str + status, stderr = self.__run_dry('{"id": "a"}\n') + self.assertEqual(status, 1) + self.assertIn("line 1", stderr) def test_extra_key_rejected(self) -> None: """Test that an extra key is rejected.""" - path: Path = self.__write_input( + status: int + status, _ = self.__run_dry( '{"id": "a", "content": "one", "extra": 1}\n') - with self.assertRaises(run_llm.InputFormatError): - run_llm.load_items(path) + self.assertEqual(status, 1) def test_non_string_content_rejected(self) -> None: """Test that a non-string content is rejected.""" - path: Path = self.__write_input('{"id": "a", "content": 3}\n') - with self.assertRaises(run_llm.InputFormatError): - run_llm.load_items(path) + status: int + status, _ = self.__run_dry('{"id": "a", "content": 3}\n') + self.assertEqual(status, 1) def test_duplicated_id_names_line(self) -> None: """Test that a duplicated ID reports the line number.""" - path: Path = self.__write_input( + status: int + stderr: str + status, stderr = self.__run_dry( '{"id": "a", "content": "one"}\n' '{"id": "a", "content": "two"}\n') - with self.assertRaises(run_llm.InputFormatError) as context: - run_llm.load_items(path) - self.assertIn("line 2", str(context.exception)) - self.assertIn("a", str(context.exception)) + self.assertEqual(status, 1) + self.assertIn("line 2", stderr) + self.assertIn("a", stderr) def test_empty_file_rejected(self) -> None: """Test that an empty input file is rejected.""" - path: Path = self.__write_input("") - with self.assertRaises(run_llm.InputFormatError): - run_llm.load_items(path) + status: int + status, _ = self.__run_dry("") + self.assertEqual(status, 1) + self.assertFalse(self.__archive_dir.exists()) class TestRequestBuilding(RunLLMTestCase): - """Test cases for the request construction.""" + """Test cases for the request preview, driven by a dry run.""" - def test_build_request(self) -> None: - """Test the shape of a batch request.""" - request: dict[str, Any] = run_llm.build_request( - run_llm.InputItem(id="song-1", content="the lyrics"), - "the system prompt", 2048, "claude-sonnet-4-6") + def setUp(self) -> None: + """Create the prompt and input files for a dry run.""" + directory: Path = self.make_temp_dir() + self.__prompt: Path = directory / "task.md" + self.__prompt.write_text( + "the system prompt", encoding="utf-8") + self.__input: Path = directory / "items.jsonl" + self.__input.write_text( + '{"id": "song-1", "content": "the lyrics"}\n', + encoding="utf-8") + self.__archive_dir: Path = directory / "runs" / "run1" + + def __preview(self, extra_argv: list[str]) -> dict[str, Any]: + """Dry-run and parse the previewed request. + + :param extra_argv: The extra command-line arguments. + :return: The parsed request. + """ + stdout: io.StringIO = io.StringIO() + with redirect_stdout(stdout): + run_llm.main([ + str(self.__prompt), str(self.__input), + str(self.__archive_dir), "--dry-run"] + extra_argv) + return json.loads(stdout.getvalue()) + + def test_default_model_request(self) -> None: + """Test the request shape for the default model.""" + request: dict[str, Any] = self.__preview([]) self.assertEqual(request["custom_id"], "song-1") params: dict[str, Any] = request["params"] self.assertEqual(params["model"], "claude-sonnet-4-6") @@ -177,125 +204,19 @@ class TestRequestBuilding(RunLLMTestCase): self.assertEqual(params["thinking"], {"type": "disabled"}) self.assertEqual(params["max_tokens"], 2048) self.assertEqual(params["system"], "the system prompt") - self.assertEqual(params["messages"], - [{"role": "user", "content": "the lyrics"}]) + self.assertEqual( + params["messages"], + [{"role": "user", "content": "the lyrics"}]) - def test_build_request_fable_5(self) -> None: + def test_fable_5_request(self) -> None: """Test the request shape for the claude-fable-5 model.""" - request: dict[str, Any] = run_llm.build_request( - run_llm.InputItem(id="group-1", content="the groups"), - "the system prompt", 8192, "claude-fable-5") + request: dict[str, Any] = self.__preview( + ["--model", "claude-fable-5", "--max-tokens", "8192"]) params: dict[str, Any] = request["params"] self.assertEqual(params["model"], "claude-fable-5") self.assertNotIn("temperature", params) self.assertNotIn("thinking", params) self.assertEqual(params["max_tokens"], 8192) - self.assertEqual(params["system"], "the system prompt") - self.assertEqual(params["messages"], - [{"role": "user", "content": "the groups"}]) - - -class TestCollectResults(RunLLMTestCase): - """Test cases for the batch result collection.""" - - def test_collect_success_and_error(self) -> None: - """Test collecting succeeded and errored results.""" - client: mock.Mock = mock.Mock() - client.messages.batches.results.return_value = iter([ - self._make_success_entry("a", "output a"), - self._make_error_entry("b", "invalid_request_error")]) - results: run_llm.Results = run_llm.collect_results( - client, "batch_x") - self.assertEqual(results["a"].text, "output a") - self.assertEqual(results["a"].stop_reason, "end_turn") - self.assertEqual(results["a"].usage, - {"input_tokens": 10, "output_tokens": 5}) - self.assertEqual(results["b"], run_llm.BatchResult( - id="b", error="invalid_request_error")) - client.messages.batches.results.assert_called_once_with( - "batch_x") - - def test_find_failures(self) -> None: - """Test finding failed and missing items.""" - results: run_llm.Results = { - "a": run_llm.BatchResult(id="a", text="fine"), - "b": run_llm.BatchResult(id="b", error="errored")} - self.assertEqual( - run_llm.find_failures(["a", "b", "c"], results), - ["b", "c"]) - - def test_sum_usage(self) -> None: - """Test summing the token usage across results.""" - results: run_llm.Results = { - "a": run_llm.BatchResult( - id="a", text="fine", - usage={"input_tokens": 10, "output_tokens": 5}), - "b": run_llm.BatchResult( - id="b", text="fine", - usage={"input_tokens": 3, "output_tokens": 2}), - "c": run_llm.BatchResult(id="c", error="errored")} - self.assertEqual( - run_llm.sum_usage(results), - {"input_tokens": 13, "output_tokens": 7}) - - -class TestArchive(RunLLMTestCase): - """Test cases for the archive directory handling.""" - - def setUp(self) -> None: - """Create a temporary directory as the runs root.""" - self.__dir: Path = self._make_temp_dir() - - def test_create_archive_dir(self) -> None: - """Test the archive directory creation.""" - target: Path = self.__dir / "01-01-tag" / "run1" - directory: Path = run_llm.create_archive_dir(target, False) - self.assertTrue(directory.is_dir()) - self.assertEqual(directory, target) - - def test_existing_archive_dir_rejected_without_replace( - self) -> None: - """Test that an existing archive is rejected by default.""" - target: Path = self.__dir / "01-01-tag" / "run1" - run_llm.create_archive_dir(target, False) - with self.assertRaises(FileExistsError): - run_llm.create_archive_dir(target, False) - - def test_existing_archive_dir_replaced(self) -> None: - """Test that --replace replaces an existing archive.""" - target: Path = self.__dir / "01-01-tag" / "run1" - first: Path = run_llm.create_archive_dir(target, False) - (first / "stale.txt").write_text("stale", encoding="utf-8") - second: Path = run_llm.create_archive_dir(target, True) - self.assertEqual(first, second) - self.assertFalse((second / "stale.txt").exists()) - - def test_replace_leaves_sibling_dir_untouched(self) -> None: - """Test that replacing run2 does not touch run1.""" - run1: Path = run_llm.create_archive_dir( - self.__dir / "01-01-tag" / "run1", False) - (run1 / "output.jsonl").write_text( - "run1 data", encoding="utf-8") - run2: Path = run_llm.create_archive_dir( - self.__dir / "01-01-tag" / "run2", False) - (run2 / "stale.jsonl").write_text("stale", encoding="utf-8") - run_llm.create_archive_dir( - self.__dir / "01-01-tag" / "run2", True) - self.assertEqual( - (run1 / "output.jsonl").read_text(encoding="utf-8"), - "run1 data") - - def test_write_jsonl(self) -> None: - """Test writing records as JSON Lines.""" - path: Path = self.__dir / "out.jsonl" - run_llm.write_jsonl(path, [{"id": "a", "text": "中文"}, - {"id": "b", "text": "two"}]) - lines: list[str] = path.read_text( - encoding="utf-8").splitlines() - self.assertEqual(len(lines), 2) - self.assertEqual(json.loads(lines[0]), - {"id": "a", "text": "中文"}) - self.assertIn("中文", lines[0]) class TestMainFlow(RunLLMTestCase): @@ -303,7 +224,7 @@ class TestMainFlow(RunLLMTestCase): def setUp(self) -> None: """Create a temporary directory with the input files.""" - directory: Path = self._make_temp_dir() + directory: Path = self.make_temp_dir() self.__runs: Path = directory / "runs" self.__archive_dir: Path = self.__runs / "task_v1" / "run1" self.__prompt: Path = directory / "task_v1.md" @@ -322,8 +243,7 @@ class TestMainFlow(RunLLMTestCase): ANTHROPIC_API_KEY="test-key") config.set_settings(self.__settings) - @staticmethod - def __make_client(entries: list[Any]) -> mock.Mock: + def __make_client(self, entries: list[Any]) -> mock.Mock: """Create a mock Anthropic client serving canned results. :param entries: The result entries of the single run batch. @@ -386,10 +306,12 @@ class TestMainFlow(RunLLMTestCase): "Done. 2 jobs finished. 02:05 elapsed.")) def test_run_produces_output_file(self) -> None: - """Test that a run submits one batch and writes output.""" + """Test that a run submits one batch and writes output, + the item order preserved and the token usage summed, the + output text written unescaped.""" client: mock.Mock = self.__make_client( - [self._make_success_entry("a", "answer a"), - self._make_success_entry("b", "answer b")]) + [self.make_success_entry("a", "answer 中文 a"), + self.make_success_entry("b", "answer b")]) status: int stderr: str with mock.patch( @@ -400,10 +322,13 @@ class TestMainFlow(RunLLMTestCase): client.messages.batches.create.call_count, 1) run_dir: Path = self.__archive_dir self.assertTrue((run_dir / "output.jsonl").exists()) + output_text: str = (run_dir / "output.jsonl").read_text( + encoding="utf-8") + self.assertIn("中文", output_text) output: list[dict[str, Any]] = [ - json.loads(x) for x in (run_dir / "output.jsonl") - .read_text(encoding="utf-8").splitlines()] - self.assertEqual(output[0]["text"], "answer a") + json.loads(x) for x in output_text.splitlines()] + self.assertEqual(output[0]["text"], "answer 中文 a") + self.assertEqual(output[1]["text"], "answer b") meta: dict[str, Any] = json.loads( (run_dir / "meta.json").read_text(encoding="utf-8")) self.assertNotIn("run", meta) @@ -431,8 +356,8 @@ class TestMainFlow(RunLLMTestCase): (run_dir / "stale.jsonl").write_text( "stale", encoding="utf-8") client: mock.Mock = self.__make_client( - [self._make_success_entry("a", "answer a"), - self._make_success_entry("b", "answer b")]) + [self.make_success_entry("a", "answer a"), + self.make_success_entry("b", "answer b")]) status: int = self.__run_main( self.__argv + ["--replace"], client)[0] self.assertEqual(status, 0) @@ -450,8 +375,8 @@ class TestMainFlow(RunLLMTestCase): (run2_dir / "output.jsonl").write_text( "stale run2 data", encoding="utf-8") client: mock.Mock = self.__make_client( - [self._make_success_entry("a", "answer a"), - self._make_success_entry("b", "answer b")]) + [self.make_success_entry("a", "answer a"), + self.make_success_entry("b", "answer b")]) argv: list[str] = [ str(self.__prompt), str(self.__input), str(run2_dir), "--replace"] @@ -465,11 +390,14 @@ class TestMainFlow(RunLLMTestCase): "stale run2 data") def test_run_failure_exits_non_zero(self) -> None: - """Test that a failed item aborts with a non-zero status.""" + """Test that a failed item aborts with a non-zero status, + the summed usage counting only the succeeded item.""" client: mock.Mock = self.__make_client( - [self._make_success_entry("a", "answer a"), - self._make_error_entry("b", "invalid_request_error")]) - status: int = self.__run_main(self.__argv, client)[0] + [self.make_success_entry("a", "answer a"), + self.make_error_entry("b", "invalid_request_error")]) + status: int + stderr: str + status, _, stderr = self.__run_main(self.__argv, client) self.assertEqual(status, 1) run_dir: Path = self.__archive_dir self.assertTrue((run_dir / "output.jsonl").exists()) @@ -479,6 +407,22 @@ class TestMainFlow(RunLLMTestCase): self.assertEqual(json.loads(output_lines[1]), {"id": "b", "error": "invalid_request_error"}) + meta: dict[str, Any] = json.loads( + (run_dir / "meta.json").read_text(encoding="utf-8")) + self.assertEqual(meta["usage"], + {"input_tokens": 10, "output_tokens": 5}) + self.assertNotIn("Done.", stderr) + + def test_missing_result_item_is_a_failure(self) -> None: + """Test that an item missing from the batch results is + reported as a failed item.""" + client: mock.Mock = self.__make_client( + [self.make_success_entry("a", "answer a")]) + status: int + stderr: str + status, _, stderr = self.__run_main(self.__argv, client) + self.assertEqual(status, 1) + self.assertIn("b", stderr) def test_invalid_input_exits_non_zero(self) -> None: """Test that an invalid input file aborts before archiving."""