From 74597112436aee0890e04ff2b0d8ae16f7ac0f02 Mon Sep 17 00:00:00 2001
From: Jens Tirsvad Nielsen <jenstirsvad@gmail.com>
Date: Tue, 13 Oct 2026 15:37:51 +0800
Subject: [PATCH] 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
---
 src/squirrel/__main__.py  | 10 +++++
 src/squirrel/census.py    | 63 +++++++++++++++++++++++++++
 src/squirrel/cli.py       | 67 +++++++++++++++++++++++++++++
 src/squirrel/constants.py |  3 ++
 tests/test_census.py      | 89 +++++++++++++++++++++++++++++++++++++++
 tests/test_cli.py         | 67 +++++++++++++++++++++++++++++
 tests/test_constants.py   |  4 ++
 tests/test_imports.py     | 40 ++++++++++++++++++
 tests/test_real_census.py | 21 +++++++++
 9 files changed, 364 insertions(+)
 create mode 100644 src/squirrel/__main__.py
 create mode 100644 src/squirrel/census.py
 create mode 100644 src/squirrel/cli.py
 create mode 100644 tests/test_census.py
 create mode 100644 tests/test_cli.py
 create mode 100644 tests/test_imports.py
 create mode 100644 tests/test_real_census.py

diff --git a/src/squirrel/__main__.py b/src/squirrel/__main__.py
new file mode 100644
index 0000000..4c97629
--- /dev/null
+++ b/src/squirrel/__main__.py
@@ -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())
diff --git a/src/squirrel/census.py b/src/squirrel/census.py
new file mode 100644
index 0000000..1c908f6
--- /dev/null
+++ b/src/squirrel/census.py
@@ -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)
diff --git a/src/squirrel/cli.py b/src/squirrel/cli.py
new file mode 100644
index 0000000..431d481
--- /dev/null
+++ b/src/squirrel/cli.py
@@ -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
diff --git a/src/squirrel/constants.py b/src/squirrel/constants.py
index 426f5ca..388ded4 100644
--- a/src/squirrel/constants.py
+++ b/src/squirrel/constants.py
@@ -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")
diff --git a/tests/test_census.py b/tests/test_census.py
new file mode 100644
index 0000000..3fa70f2
--- /dev/null
+++ b/tests/test_census.py
@@ -0,0 +1,89 @@
+"""Tests for the squirrel.census module."""
+
+import pathlib
+
+from squirrel import census
+from squirrel.constants import (
+    FILE_ENCODING,
+    FUR_COLOR_COLUMN,
+    FUR_COLORS,
+    SUMMARY_COLUMNS,
+)
+
+OTHER_COLUMN = "Unique Squirrel ID"
+
+
+def _row(color: str) -> dict[str, str]:
+    return {FUR_COLOR_COLUMN: color}
+
+
+def test_read_csv_returns_one_dict_per_row(tmp_path: pathlib.Path) -> None:
+    path = tmp_path / "census.csv"
+    path.write_text(
+        f"{OTHER_COLUMN},{FUR_COLOR_COLUMN}\n1,{FUR_COLORS[0]}\n2,\n",
+        encoding=FILE_ENCODING,
+    )
+
+    rows = census.read_csv(path)
+
+    assert rows == [
+        {OTHER_COLUMN: "1", FUR_COLOR_COLUMN: FUR_COLORS[0]},
+        {OTHER_COLUMN: "2", FUR_COLOR_COLUMN: ""},
+    ]
+
+
+def test_count_by_fur_color_counts_each_colour() -> None:
+    gray, cinnamon, black = FUR_COLORS
+    rows = [
+        _row(gray),
+        _row(gray),
+        _row(cinnamon),
+        _row(black),
+        _row(black),
+        _row(black),
+    ]
+
+    counts = census.count_by_fur_color(rows)
+
+    assert counts == {gray: 2, cinnamon: 1, black: 3}
+
+
+def test_count_by_fur_color_ignores_blank_and_unknown_colours() -> None:
+    rows = [_row(""), _row("White"), _row(FUR_COLORS[0].lower()), {OTHER_COLUMN: "1"}]
+
+    counts = census.count_by_fur_color(rows)
+
+    assert sum(counts.values()) == 0
+
+
+def test_count_by_fur_color_includes_zero_counts_in_colour_order() -> None:
+    counts = census.count_by_fur_color([_row(FUR_COLORS[1])])
+
+    assert list(counts.items()) == [
+        (FUR_COLORS[0], 0),
+        (FUR_COLORS[1], 1),
+        (FUR_COLORS[2], 0),
+    ]
+
+
+def test_count_by_fur_color_uses_given_colours() -> None:
+    counts = census.count_by_fur_color([_row("White"), _row("White")], colors=["White"])
+
+    assert counts == {"White": 2}
+
+
+def test_build_summary_keeps_order_of_counts() -> None:
+    counts = {"Black": 3, "Gray": 1, "Cinnamon": 2}
+
+    assert census.build_summary(counts) == [("Black", 3), ("Gray", 1), ("Cinnamon", 2)]
+
+
+def test_write_csv_writes_header_then_rows(tmp_path: pathlib.Path) -> None:
+    path = tmp_path / "summary.csv"
+    summary = [(FUR_COLORS[0], 5), (FUR_COLORS[1], 0)]
+
+    census.write_csv(summary, path)
+
+    header = ",".join(SUMMARY_COLUMNS)
+    expected = f"{header}\n{FUR_COLORS[0]},5\n{FUR_COLORS[1]},0\n"
+    assert path.read_bytes().decode(FILE_ENCODING).replace("\r\n", "\n") == expected
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..7812a43
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,67 @@
+"""Tests for the squirrel.cli module."""
+
+import pathlib
+
+import pytest
+
+from squirrel import cli
+from squirrel.constants import (
+    FILE_ENCODING,
+    FUR_COLOR_COLUMN,
+    FUR_COLORS,
+    SUMMARY_COLUMNS,
+)
+
+
+def _write_census(path: pathlib.Path) -> None:
+    gray, cinnamon, _black = FUR_COLORS
+    lines = [FUR_COLOR_COLUMN, gray, gray, cinnamon, ""]
+    path.write_text("\n".join(lines) + "\n", encoding=FILE_ENCODING)
+
+
+def test_main_writes_summary_file_and_returns_zero(tmp_path: pathlib.Path) -> None:
+    census_path = tmp_path / "census.csv"
+    summary_path = tmp_path / "summary.csv"
+    _write_census(census_path)
+
+    status = cli.main([str(census_path), str(summary_path)])
+
+    assert status == 0
+    lines = summary_path.read_text(encoding=FILE_ENCODING).splitlines()
+    assert lines == [
+        ",".join(SUMMARY_COLUMNS),
+        f"{FUR_COLORS[0]},2",
+        f"{FUR_COLORS[1]},1",
+        f"{FUR_COLORS[2]},0",
+    ]
+
+
+def test_main_prints_one_line_per_colour(
+    tmp_path: pathlib.Path, capsys: pytest.CaptureFixture[str]
+) -> None:
+    census_path = tmp_path / "census.csv"
+    _write_census(census_path)
+
+    cli.main([str(census_path), str(tmp_path / "summary.csv")])
+
+    out_lines = capsys.readouterr().out.splitlines()
+    assert out_lines == [
+        f"{FUR_COLORS[0]}: 2",
+        f"{FUR_COLORS[1]}: 1",
+        f"{FUR_COLORS[2]}: 0",
+    ]
+
+
+def test_main_returns_2_and_writes_nothing_when_census_missing(
+    tmp_path: pathlib.Path, capsys: pytest.CaptureFixture[str]
+) -> None:
+    missing = tmp_path / "missing.csv"
+    summary_path = tmp_path / "summary.csv"
+
+    status = cli.main([str(missing), str(summary_path)])
+
+    captured = capsys.readouterr()
+    assert status == 2
+    assert not summary_path.exists()
+    assert captured.err.strip() == f"Census file not found: {missing}"
+    assert captured.out == ""
diff --git a/tests/test_constants.py b/tests/test_constants.py
index 35527f7..3af4440 100644
--- a/tests/test_constants.py
+++ b/tests/test_constants.py
@@ -31,6 +31,10 @@ def test_file_encoding_and_column_name_values() -> None:
     assert constants.FUR_COLOR_COLUMN == "Primary Fur Color"
 
 
+def test_exit_input_missing_is_2() -> None:
+    assert constants.EXIT_INPUT_MISSING == 2
+
+
 def test_module_imports_only_standard_library() -> None:
     source = pathlib.Path(constants.__file__).read_text(encoding="utf-8")
     imported: set[str] = set()
diff --git a/tests/test_imports.py b/tests/test_imports.py
new file mode 100644
index 0000000..69d369c
--- /dev/null
+++ b/tests/test_imports.py
@@ -0,0 +1,40 @@
+"""Every module of the package imports only the standard library or itself."""
+
+import ast
+import pathlib
+import sys
+
+import squirrel
+
+PACKAGE_NAME = squirrel.__name__
+PACKAGE_DIR = pathlib.Path(squirrel.__file__).parent
+ALLOWED = sys.stdlib_module_names | {PACKAGE_NAME}
+
+
+def _imported_top_level_names(source: str) -> set[str]:
+    imported: set[str] = set()
+    for node in ast.walk(ast.parse(source)):
+        if isinstance(node, ast.Import):
+            imported.update(alias.name.split(".")[0] for alias in node.names)
+        elif (
+            isinstance(node, ast.ImportFrom)
+            and node.module is not None
+            and node.level == 0
+        ):
+            imported.add(node.module.split(".")[0])
+    return imported
+
+
+def test_package_has_modules_to_check() -> None:
+    assert list(PACKAGE_DIR.glob("*.py"))
+
+
+def test_every_module_imports_only_standard_library_or_squirrel() -> None:
+    offenders: dict[str, set[str]] = {}
+    for module in PACKAGE_DIR.rglob("*.py"):
+        source = module.read_text(encoding="utf-8")
+        foreign = _imported_top_level_names(source) - ALLOWED
+        if foreign:
+            offenders[module.name] = foreign
+
+    assert offenders == {}
diff --git a/tests/test_real_census.py b/tests/test_real_census.py
new file mode 100644
index 0000000..6b09145
--- /dev/null
+++ b/tests/test_real_census.py
@@ -0,0 +1,21 @@
+"""Check the counts of the real 2018 Central Park census file, when present."""
+
+import pathlib
+
+import pytest
+
+from squirrel import census
+from squirrel.constants import DEFAULT_INPUT_PATH, FUR_COLORS
+
+REPO_ROOT = pathlib.Path(__file__).parents[1]
+
+
+def test_real_census_counts_match_published_figures() -> None:
+    path = REPO_ROOT / DEFAULT_INPUT_PATH
+    if not path.is_file():
+        pytest.skip(f"census file not present: {path}")
+
+    counts = census.count_by_fur_color(census.read_csv(path))
+
+    gray, cinnamon, black = FUR_COLORS
+    assert counts == {gray: 2473, cinnamon: 392, black: 103}