Implement census analysis: read, count, summarise, write and CLI

Closes #11
Closes #12
Closes #13
Closes #14
Closes #15

Task: MIL-002#4
Task: MIL-002#5
Task: MIL-002#6
Task: MIL-002#7
Task: MIL-002#8
This commit is contained in:
2026-10-13 15:37:51 +08:00
parent 259d017ad2
commit 7459711243
9 changed files with 364 additions and 0 deletions
+10
View File
@@ -0,0 +1,10 @@
##
# @file __main__.py
# @brief Allows `python -m squirrel` to run the command line.
"""Allows `python -m squirrel` to run the command line."""
from squirrel.cli import main
if __name__ == "__main__":
raise SystemExit(main())
+63
View File
@@ -0,0 +1,63 @@
##
# @file census.py
# @brief Pure steps of the census analysis: read, count, summarise, write.
"""Pure steps of the census analysis: read, count, summarise, write."""
import collections
import csv
import pathlib
from collections.abc import Iterable, Mapping, Sequence
from squirrel.constants import (
FILE_ENCODING,
FUR_COLOR_COLUMN,
FUR_COLORS,
SUMMARY_COLUMNS,
)
##
# @brief Read a CSV file into one dictionary per row, keyed by the header names.
# @param path Location of the CSV file.
# @return The rows of the file, in file order.
def read_csv(path: pathlib.Path) -> list[dict[str, str]]:
"""Read a CSV file into one dictionary per row, keyed by the header names."""
with path.open(encoding=FILE_ENCODING, newline="") as file:
return list(csv.DictReader(file))
##
# @brief Count the rows per fur colour; other or blank colours are counted in none.
# @param rows Census rows, each holding the fur colour column.
# @param colors Fur colours to count, in the order of the result.
# @return Count per colour of `colors`, 0 for a colour that does not occur.
def count_by_fur_color(
rows: Iterable[Mapping[str, str]],
colors: Sequence[str] = FUR_COLORS,
) -> dict[str, int]:
"""Count the rows per fur colour; other or blank colours are counted in none."""
seen = collections.Counter(row.get(FUR_COLOR_COLUMN) for row in rows)
return {color: seen[color] for color in colors}
##
# @brief Turn the counts into summary rows, keeping the order of the counts.
# @param counts Count per fur colour.
# @return One (colour, count) pair per colour.
def build_summary(counts: Mapping[str, int]) -> list[tuple[str, int]]:
"""Turn the counts into summary rows, keeping the order of the counts."""
return list(counts.items())
##
# @brief Write the summary as a CSV file: the header row, then one row per colour.
# @param summary Rows to write, as built by `build_summary`.
# @param path Location of the CSV file; its folder must exist.
# @return None.
def write_csv(summary: Sequence[tuple[str, int]], path: pathlib.Path) -> None:
"""Write the summary as a CSV file: the header row, then one row per colour."""
with path.open("w", encoding=FILE_ENCODING, newline="") as file:
writer = csv.writer(file)
writer.writerow(SUMMARY_COLUMNS)
writer.writerows(summary)
+67
View File
@@ -0,0 +1,67 @@
##
# @file cli.py
# @brief Command line entry point: count the squirrels and write the summary.
"""Command line entry point: count the squirrels and write the summary."""
import argparse
import pathlib
import sys
from collections.abc import Sequence
from squirrel import census
from squirrel.constants import (
DEFAULT_INPUT_PATH,
DEFAULT_OUTPUT_PATH,
EXIT_INPUT_MISSING,
)
##
# @brief Build the argument parser of the command line.
# @return The parser with the optional `input` and `output` positionals.
def _build_parser() -> argparse.ArgumentParser:
"""Build the argument parser of the command line."""
parser = argparse.ArgumentParser(
prog="squirrel",
description="Count the census squirrels by primary fur colour.",
)
parser.add_argument(
"input",
nargs="?",
type=pathlib.Path,
default=DEFAULT_INPUT_PATH,
help="census CSV file (default: %(default)s)",
)
parser.add_argument(
"output",
nargs="?",
type=pathlib.Path,
default=DEFAULT_OUTPUT_PATH,
help="summary CSV file to write (default: %(default)s)",
)
return parser
##
# @brief Run the census analysis from the command line.
# @param argv Arguments without the program name; `sys.argv` is used when None.
# @return 0 on success, 2 when the census file does not exist.
def main(argv: Sequence[str] | None = None) -> int:
"""Run the census analysis from the command line."""
args = _build_parser().parse_args(argv)
input_path: pathlib.Path = args.input
output_path: pathlib.Path = args.output
if not input_path.is_file():
print(f"Census file not found: {input_path}", file=sys.stderr)
return EXIT_INPUT_MISSING
counts = census.count_by_fur_color(census.read_csv(input_path))
summary = census.build_summary(counts)
census.write_csv(summary, output_path)
# print is the intended output here: the counts are shown to the participant.
for color, count in summary:
print(f"{color}: {count}")
return 0
+3
View File
@@ -23,5 +23,8 @@ FUR_COLOR_COLUMN: str = "Primary Fur Color"
## @brief Fur colours that are counted, in the order they appear in the summary.
FUR_COLORS: tuple[str, ...] = ("Gray", "Cinnamon", "Black")
## @brief Exit status of the command line when the census file does not exist.
EXIT_INPUT_MISSING: int = 2
## @brief Header of the summary file: the colour column, then the count column.
SUMMARY_COLUMNS: tuple[str, str] = ("Fur Color", "Count")