From 5617d9c03eca4ff9d2332de47f803ad948a72715 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Fri, 10 Jul 2026 22:47:17 +0200 Subject: [PATCH 01/11] Adds mechanism to compute statistics and metadata for each snapshot given a reader. --- src/dstrack/snapshot/__init__.py | 39 +++ src/dstrack/snapshot/_metadata.py | 134 +++++++++ src/dstrack/snapshot/_stats.py | 435 ++++++++++++++++++++++++++++++ src/dstrack/snapshot/_version.py | 12 + tests/test_snapshot_metadata.py | 193 +++++++++++++ tests/test_snapshot_stats.py | 364 +++++++++++++++++++++++++ 6 files changed, 1177 insertions(+) create mode 100644 src/dstrack/snapshot/__init__.py create mode 100644 src/dstrack/snapshot/_metadata.py create mode 100644 src/dstrack/snapshot/_stats.py create mode 100644 src/dstrack/snapshot/_version.py create mode 100644 tests/test_snapshot_metadata.py create mode 100644 tests/test_snapshot_stats.py diff --git a/src/dstrack/snapshot/__init__.py b/src/dstrack/snapshot/__init__.py new file mode 100644 index 0000000..9a2e9f6 --- /dev/null +++ b/src/dstrack/snapshot/__init__.py @@ -0,0 +1,39 @@ +"""Snapshot computation for dstrack. + +Classes +------- +:class:`MetadataBuilder` + Builds identity and structural snapshot fields from a reader's schema. + +:class:`StatsComputer` + Computes per-column and dataset-level statistics in a single data pass. + +Result types +------------ +:class:`SnapshotMetadata`, :class:`DatasetStats` +""" + +from ._metadata import MetadataBuilder, SnapshotMetadata +from ._stats import ( + DatasetStats, + DatetimeColumnStats, + HistogramStats, + NumericColumnStats, + OtherColumnStats, + PercentileStats, + StatsComputer, + StringColumnStats, +) + +__all__ = [ + "DatasetStats", + "DatetimeColumnStats", + "HistogramStats", + "MetadataBuilder", + "NumericColumnStats", + "OtherColumnStats", + "PercentileStats", + "SnapshotMetadata", + "StatsComputer", + "StringColumnStats", +] diff --git a/src/dstrack/snapshot/_metadata.py b/src/dstrack/snapshot/_metadata.py new file mode 100644 index 0000000..8e56414 --- /dev/null +++ b/src/dstrack/snapshot/_metadata.py @@ -0,0 +1,134 @@ +import hashlib +import uuid +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from dstrack.readers import ColumnInfo, TabularReader +from dstrack.snapshot._version import FORMAT_VERSION + + +@dataclass +class SnapshotMetadata: + """Identity and structural fields for a snapshot. + + Covers every required top-level field that does not require a data pass: + versioning identifiers, authorship, schema shape, and the source hash. + """ + + format_version: str + snapshot_id: str + created_at: str + created_by: str + dataset_name: str + dataset_path: str + source_type: str + source_hash: str + num_columns: int + columns: list[dict[str, Any]] + schema_hash: str + + +class MetadataBuilder: + """Builds identity and structural metadata without reading data rows. + + Computes snapshot_id, created_at, created_by, dataset_name, + dataset_path, source_type, source_hash, num_columns, columns, and + schema_hash from the reader's column descriptors alone. + + Note: + schema_hash is order-independent: reordering a dataset's columns + without changing their names or dtypes produces the same hash. + """ + + def build( + self, + reader: TabularReader, + *, + dataset_name: str, + dataset_path: str | Path, + source_type: str, + created_by: str, + source_hash: str | None = None, + ) -> SnapshotMetadata: + """Build metadata for a snapshot. + + Args: + reader: Any TabularReader; only ``columns()`` is called. + dataset_name: Human-readable dataset name stored in the snapshot. + dataset_path: Source path or URI at snapshot time. + source_type: Origin kind (``"file"``, ``"directory"``, etc.). + created_by: User or process identifier. + source_hash: Pre-computed source hash. When ``None`` and + ``dataset_path`` resolves to a regular file, a SHA-256 of + the file bytes is computed automatically. + + Returns: + A populated :class:`SnapshotMetadata` instance. + """ + cols = reader.columns() + path_str = str(dataset_path) + + if source_hash is None: + p = Path(dataset_path) + source_hash = _hash_file(p) if p.is_file() else "" + + return SnapshotMetadata( + format_version=FORMAT_VERSION, + snapshot_id=str(uuid.uuid4()), + created_at=datetime.now(UTC).isoformat(), + created_by=created_by, + dataset_name=dataset_name, + dataset_path=path_str, + source_type=source_type, + source_hash=source_hash, + num_columns=len(cols), + columns=_columns_to_dicts(cols), + schema_hash=_schema_hash(cols), + ) + + +def _hash_file(path: Path) -> str: + """Compute the SHA-256 hex digest of a file's contents. + + Args: + path: Path to the file to hash. + + Returns: + The hex-encoded SHA-256 digest of the file's bytes. + """ + h = hashlib.sha256() + with path.open("rb") as fh: + for chunk in iter(lambda: fh.read(65536), b""): + h.update(chunk) + return h.hexdigest() + + +def _columns_to_dicts(cols: list[ColumnInfo]) -> list[dict[str, Any]]: + """Convert column descriptors into plain dicts for serialization. + + Args: + cols: Column descriptors as returned by a reader. + + Returns: + One dict per column with ``name``, ``dtype``, and ``nullable`` keys. + """ + return [{"name": c.name, "dtype": c.dtype, "nullable": c.nullable} for c in cols] + + +def _schema_hash(cols: list[ColumnInfo]) -> str: + """Compute a deterministic hash of a schema's column names and dtypes. + + Args: + cols: Column descriptors, in any order. + + Returns: + A hex-encoded SHA-256 digest that changes if any column's name or + dtype changes. Order-independent: reordering ``cols`` does not + change the result. Ignores ``nullable``. + """ + h = hashlib.sha256() + for name, dtype in sorted((c.name, c.dtype) for c in cols): + h.update(f"{name}\x00{dtype}\n".encode()) + return h.hexdigest() diff --git a/src/dstrack/snapshot/_stats.py b/src/dstrack/snapshot/_stats.py new file mode 100644 index 0000000..6af0a4e --- /dev/null +++ b/src/dstrack/snapshot/_stats.py @@ -0,0 +1,435 @@ +import math +from collections import Counter +from dataclasses import dataclass, field +from typing import Protocol + +from dstrack.readers import Cell, TabularReader + +_NUMERIC_DTYPES: frozenset[str] = frozenset( + {"int8", "int16", "int32", "int64", "float16", "float32", "float64", "bool"} +) +_NUM_HISTOGRAM_BINS: int = 20 +_TOP_VALUES_K: int = 50 + + +@dataclass(frozen=True, slots=True) +class PercentileStats: + """p5-p99 percentiles of a numeric column, linearly interpolated.""" + + p5: float = 0.0 + p25: float = 0.0 + p50: float = 0.0 + p75: float = 0.0 + p95: float = 0.0 + p99: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class HistogramStats: + """Equal-width histogram of a numeric column's values.""" + + bin_edges: list[float] = field(default_factory=list) + counts: list[int] = field(default_factory=list) + + +@dataclass(frozen=True, slots=True) +class NumericColumnStats: + """Statistics for an ``int*``/``float*``/``bool`` column.""" + + null_count: int = 0 + null_fraction: float = 0.0 + mean: float = 0.0 + std: float = 0.0 + min: float = 0.0 + max: float = 0.0 + percentiles: PercentileStats = field(default_factory=PercentileStats) + histogram: HistogramStats = field(default_factory=HistogramStats) + num_unique: int = 0 + + +@dataclass(frozen=True, slots=True) +class StringColumnStats: + """Statistics for a ``string`` column.""" + + null_count: int = 0 + null_fraction: float = 0.0 + num_unique: int = 0 + top_values: dict[str, int] = field(default_factory=dict) + top_values_coverage: float = 0.0 + avg_char_length: float = 0.0 + min_char_length: float = 0.0 + max_char_length: float = 0.0 + avg_token_count: float = 0.0 + min_token_count: float = 0.0 + max_token_count: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class DatetimeColumnStats: + """Statistics for a ``datetime64`` column; min/max are ISO strings.""" + + null_count: int = 0 + null_fraction: float = 0.0 + min: str = "" + max: str = "" + range_days: float = 0.0 + + +@dataclass(frozen=True, slots=True) +class OtherColumnStats: + """Null-only statistics for a column whose dtype is unrecognized.""" + + null_count: int = 0 + null_fraction: float = 0.0 + + +# The per-column entry type stored in DatasetStats.column_stats, keyed by dtype category. +ColumnStats = ( + NumericColumnStats | StringColumnStats | DatetimeColumnStats | OtherColumnStats +) + + +@dataclass +class DatasetStats: + """Per-column and dataset-level statistics produced by a data pass. + + Covers num_rows, column_stats, duplicate_row_fraction, constant_columns, + and high_null_columns. Sketch-based fields (near_duplicate_estimate, + row_minhash, row_hyperloglog) are handled by a separate builder. + """ + + num_rows: int + column_stats: dict[str, ColumnStats] + duplicate_row_fraction: float + constant_columns: list[str] + high_null_columns: list[str] + + +class _ColumnAcc(Protocol): + """Common interface every per-column accumulator implements. + + ``compute`` drives all accumulators through this interface so it does not + need to branch on column type while scanning rows or building stats. + """ + + null_count: int + + def update(self, val: Cell) -> None: ... + + def is_constant(self) -> bool: ... + + def stats(self, null_fraction: float) -> ColumnStats: ... + + +@dataclass +class _NumericAcc: + """Running accumulator for an ``int*``/``float*``/``bool`` column.""" + + vals: list[float] = field(default_factory=list) + null_count: int = 0 + + def update(self, val: Cell) -> None: + self.vals.append(float(val)) # type: ignore[arg-type] + + def is_constant(self) -> bool: + return not self.vals or min(self.vals) == max(self.vals) + + def stats(self, null_fraction: float) -> NumericColumnStats: + return _numeric_stats(self.vals, self.null_count, null_fraction) + + +@dataclass +class _StringAcc: + """Running accumulator for a ``string`` column.""" + + counter: Counter[str] = field(default_factory=Counter) + char_len_sum: int = 0 + char_len_min: float = math.inf + char_len_max: float = -math.inf + token_count_sum: int = 0 + token_count_min: float = math.inf + token_count_max: float = -math.inf + non_null_count: int = 0 + null_count: int = 0 + + def update(self, val: Cell) -> None: + s = str(val) + self.non_null_count += 1 + self.counter[s] += 1 + cl = len(s) + tc = len(s.split()) + self.char_len_sum += cl + self.token_count_sum += tc + self.char_len_min = min(self.char_len_min, cl) + self.char_len_max = max(self.char_len_max, cl) + self.token_count_min = min(self.token_count_min, tc) + self.token_count_max = max(self.token_count_max, tc) + + def is_constant(self) -> bool: + return len(self.counter) <= 1 + + def stats(self, null_fraction: float) -> StringColumnStats: + return _string_stats(self, null_fraction) + + +@dataclass +class _DatetimeAcc: + """Running accumulator for a ``datetime64`` column. + + Values are compared and stored as ISO strings rather than parsed + ``datetime`` objects, since lexicographic order matches chronological + order for ISO 8601 timestamps. + """ + + min_val: str = "" + max_val: str = "" + non_null_count: int = 0 + null_count: int = 0 + + def update(self, val: Cell) -> None: + s = str(val) + if self.non_null_count == 0 or s < self.min_val: + self.min_val = s + if self.non_null_count == 0 or s > self.max_val: + self.max_val = s + self.non_null_count += 1 + + def is_constant(self) -> bool: + return self.non_null_count == 0 or self.min_val == self.max_val + + def stats(self, null_fraction: float) -> DatetimeColumnStats: + return _datetime_stats(self, null_fraction) + + +@dataclass +class _OtherAcc: + """Accumulator for columns whose dtype is unrecognized. + + Only null counts are tracked; the value itself is never inspected. + """ + + null_count: int = 0 + + def update(self, val: Cell) -> None: + pass + + def is_constant(self) -> bool: + return False + + def stats(self, null_fraction: float) -> OtherColumnStats: + return OtherColumnStats(null_count=self.null_count, null_fraction=null_fraction) + + +class StatsComputer: + """Computes per-column and dataset-level statistics in a single data pass. + + Handles ``int*``, ``float*``, and ``bool`` dtypes as numeric, + ``string`` columns, and ``datetime64`` columns. Unknown dtypes produce + only null counts. + """ + + def compute(self, reader: TabularReader) -> DatasetStats: + """Run a full data pass and return aggregated statistics. + + Args: + reader: Any TabularReader whose batches will be consumed once. + + Returns: + A populated :class:`DatasetStats` instance. + """ + cols = reader.columns() + col_names = [c.name for c in cols] + col_dtypes = {c.name: c.dtype for c in cols} + + accs = self._init_accumulators(col_dtypes) + + num_rows = 0 + seen_rows: set[tuple[Cell, ...]] = set() + duplicate_count = 0 + + for batch in reader.iter_batches(): + batch_rows, batch_duplicates = self._process_batch( + batch, col_names, accs, seen_rows + ) + num_rows += batch_rows + duplicate_count += batch_duplicates + + return self._build_dataset_stats(col_names, accs, num_rows, duplicate_count) + + def _init_accumulators(self, col_dtypes: dict[str, str]) -> dict[str, _ColumnAcc]: + """Create one accumulator per column, chosen by dtype.""" + accs: dict[str, _ColumnAcc] = {} + for name, dtype in col_dtypes.items(): + if dtype in _NUMERIC_DTYPES: + accs[name] = _NumericAcc() + elif dtype == "string": + accs[name] = _StringAcc() + elif dtype == "datetime64": + accs[name] = _DatetimeAcc() + else: + accs[name] = _OtherAcc() + return accs + + def _process_batch( + self, + batch: list[list[Cell]], + col_names: list[str], + accs: dict[str, _ColumnAcc], + seen_rows: set[tuple[Cell, ...]], + ) -> tuple[int, int]: + """Feed one batch of rows into the accumulators. + + Returns the (row_count, duplicate_count) contributed by this batch; + ``seen_rows`` is shared and updated across batches so duplicates are + detected dataset-wide, not just within a single batch. + """ + num_rows = 0 + duplicate_count = 0 + for row in batch: + num_rows += 1 + row_key = tuple(row) + if row_key in seen_rows: + duplicate_count += 1 + else: + seen_rows.add(row_key) + + for name, val in zip(col_names, row, strict=True): + acc = accs[name] + if val is None: + acc.null_count += 1 + else: + acc.update(val) + return num_rows, duplicate_count + + def _build_dataset_stats( + self, + col_names: list[str], + accs: dict[str, _ColumnAcc], + num_rows: int, + duplicate_count: int, + ) -> DatasetStats: + """Turn the fully-populated accumulators into a DatasetStats.""" + column_stats: dict[str, ColumnStats] = {} + constant_columns: list[str] = [] + high_null_columns: list[str] = [] + + for name in col_names: + acc = accs[name] + null_fraction = acc.null_count / num_rows if num_rows > 0 else 0.0 + if null_fraction > 0.5: + high_null_columns.append(name) + column_stats[name] = acc.stats(null_fraction) + if acc.is_constant(): + constant_columns.append(name) + + return DatasetStats( + num_rows=num_rows, + column_stats=column_stats, + duplicate_row_fraction=duplicate_count / num_rows if num_rows > 0 else 0.0, + constant_columns=constant_columns, + high_null_columns=high_null_columns, + ) + + +def _numeric_stats( + vals: list[float], null_count: int, null_fraction: float +) -> NumericColumnStats: + """Build NumericColumnStats from the raw values collected by a _NumericAcc.""" + n = len(vals) + if n == 0: + return NumericColumnStats(null_count=null_count, null_fraction=null_fraction) + + sorted_vals = sorted(vals) + mean = sum(vals) / n + variance = sum((v - mean) ** 2 for v in vals) / n + return NumericColumnStats( + null_count=null_count, + null_fraction=null_fraction, + mean=mean, + std=math.sqrt(variance), + min=sorted_vals[0], + max=sorted_vals[-1], + percentiles=PercentileStats( + p5=_percentile(sorted_vals, 5), + p25=_percentile(sorted_vals, 25), + p50=_percentile(sorted_vals, 50), + p75=_percentile(sorted_vals, 75), + p95=_percentile(sorted_vals, 95), + p99=_percentile(sorted_vals, 99), + ), + histogram=_histogram(sorted_vals, _NUM_HISTOGRAM_BINS), + num_unique=len(set(vals)), + ) + + +def _percentile(sorted_vals: list[float], p: int) -> float: + """Linearly interpolated p-th percentile (0-100) of already-sorted values.""" + n = len(sorted_vals) + if n == 1: + return sorted_vals[0] + idx = (p / 100) * (n - 1) + lo = int(idx) + hi = lo + 1 + if hi >= n: + return sorted_vals[-1] + return sorted_vals[lo] + (idx - lo) * (sorted_vals[hi] - sorted_vals[lo]) + + +def _histogram(sorted_vals: list[float], num_bins: int) -> HistogramStats: + """Bucket already-sorted values into num_bins equal-width bins.""" + lo, hi = sorted_vals[0], sorted_vals[-1] + if lo == hi: + return HistogramStats(bin_edges=[lo, lo + 1.0], counts=[len(sorted_vals)]) + step = (hi - lo) / num_bins + edges = [lo + i * step for i in range(num_bins + 1)] + counts = [0] * num_bins + for v in sorted_vals: + idx = min(int((v - lo) / step), num_bins - 1) + counts[idx] += 1 + return HistogramStats(bin_edges=edges, counts=counts) + + +def _string_stats(acc: _StringAcc, null_fraction: float) -> StringColumnStats: + """Build StringColumnStats from a fully-populated _StringAcc.""" + n = acc.non_null_count + top_k = acc.counter.most_common(_TOP_VALUES_K) + top_values_coverage = sum(c for _, c in top_k) / n if n > 0 else 0.0 + return StringColumnStats( + null_count=acc.null_count, + null_fraction=null_fraction, + num_unique=len(acc.counter), + top_values=dict(top_k), + top_values_coverage=top_values_coverage, + avg_char_length=acc.char_len_sum / n if n > 0 else 0.0, + min_char_length=acc.char_len_min if n > 0 else 0.0, + max_char_length=acc.char_len_max if n > 0 else 0.0, + avg_token_count=acc.token_count_sum / n if n > 0 else 0.0, + min_token_count=acc.token_count_min if n > 0 else 0.0, + max_token_count=acc.token_count_max if n > 0 else 0.0, + ) + + +def _datetime_stats(acc: _DatetimeAcc, null_fraction: float) -> DatetimeColumnStats: + """Build DatetimeColumnStats from a fully-populated _DatetimeAcc.""" + if acc.non_null_count == 0: + return DatetimeColumnStats( + null_count=acc.null_count, null_fraction=null_fraction + ) + + from datetime import datetime as _dt + + range_days = 0.0 + try: + min_dt = _dt.fromisoformat(acc.min_val.replace("Z", "+00:00")) + max_dt = _dt.fromisoformat(acc.max_val.replace("Z", "+00:00")) + range_days = (max_dt - min_dt).total_seconds() / 86400 + except ValueError: + pass + + return DatetimeColumnStats( + null_count=acc.null_count, + null_fraction=null_fraction, + min=acc.min_val, + max=acc.max_val, + range_days=range_days, + ) diff --git a/src/dstrack/snapshot/_version.py b/src/dstrack/snapshot/_version.py new file mode 100644 index 0000000..2a7a255 --- /dev/null +++ b/src/dstrack/snapshot/_version.py @@ -0,0 +1,12 @@ +"""Single source of truth for the snapshot format version. + +Bump ``FORMAT_VERSION`` and ``SCHEMA_PATH`` together whenever +``SnapshotMetadata``'s shape changes. +""" + +from pathlib import Path + +FORMAT_VERSION = "1" +SCHEMA_PATH = ( + Path(__file__).parent.parent / "schemas" / f"snapshot_v{FORMAT_VERSION}.json" +) diff --git a/tests/test_snapshot_metadata.py b/tests/test_snapshot_metadata.py new file mode 100644 index 0000000..d6ac741 --- /dev/null +++ b/tests/test_snapshot_metadata.py @@ -0,0 +1,193 @@ +"""Tests for MetadataBuilder.""" + +import hashlib +import uuid +from collections.abc import Iterator +from pathlib import Path + +from dstrack.readers import ColumnInfo, TabularReader +from dstrack.snapshot import MetadataBuilder, SnapshotMetadata + +# --------------------------------------------------------------------------- +# Stubs +# --------------------------------------------------------------------------- + + +def _make_reader(*cols: ColumnInfo) -> TabularReader: + class _Stub: + def columns(self) -> list[ColumnInfo]: + return list(cols) + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[object]]]: + return iter([]) + + return _Stub() # type: ignore[return-value] + + +def _build( + reader: TabularReader, + *, + path: str | Path = "/data/ds.csv", + source_hash: str | None = None, +) -> SnapshotMetadata: + return MetadataBuilder().build( + reader, + dataset_name="my_dataset", + dataset_path=path, + source_type="file", + created_by="tester", + source_hash=source_hash, + ) + + +# --------------------------------------------------------------------------- +# Basic fields +# --------------------------------------------------------------------------- + + +def test_format_version_is_one() -> None: + """Built metadata always reports format_version '1'.""" + meta = _build(_make_reader()) + assert meta.format_version == "1" + + +def test_snapshot_id_is_valid_uuid() -> None: + """snapshot_id is a valid UUIDv4 string.""" + meta = _build(_make_reader()) + parsed = uuid.UUID(meta.snapshot_id) + assert parsed.version == 4 + + +def test_created_at_is_iso8601() -> None: + """created_at parses as a timezone-aware ISO 8601 timestamp.""" + from datetime import datetime + + meta = _build(_make_reader()) + dt = datetime.fromisoformat(meta.created_at) + assert dt.tzinfo is not None + + +def test_created_by_and_name() -> None: + """created_by, dataset_name, and source_type are passed through as given.""" + meta = _build(_make_reader()) + assert meta.created_by == "tester" + assert meta.dataset_name == "my_dataset" + assert meta.source_type == "file" + + +def test_dataset_path_stored_as_string() -> None: + """A Path dataset_path is coerced to a string on the built metadata.""" + meta = _build(_make_reader(), path=Path("/some/path/data.csv")) + assert meta.dataset_path == "/some/path/data.csv" + + +# --------------------------------------------------------------------------- +# Column descriptors +# --------------------------------------------------------------------------- + + +def test_num_columns_matches_reader() -> None: + """num_columns equals the number of columns reported by the reader.""" + cols = [ColumnInfo("a", "int64"), ColumnInfo("b", "string")] + meta = _build(_make_reader(*cols)) + assert meta.num_columns == 2 + + +def test_columns_serialised_correctly() -> None: + """Each ColumnInfo is serialised to a name/dtype/nullable dict.""" + cols = [ + ColumnInfo("x", "float64", nullable=False), + ColumnInfo("y", "string", nullable=True), + ] + meta = _build(_make_reader(*cols)) + assert meta.columns == [ + {"name": "x", "dtype": "float64", "nullable": False}, + {"name": "y", "dtype": "string", "nullable": True}, + ] + + +def test_empty_reader_produces_zero_columns() -> None: + """A reader with no columns yields num_columns 0 and an empty columns list.""" + meta = _build(_make_reader()) + assert meta.num_columns == 0 + assert meta.columns == [] + + +# --------------------------------------------------------------------------- +# schema_hash +# --------------------------------------------------------------------------- + + +def test_schema_hash_is_hex_string() -> None: + """schema_hash is a valid hexadecimal string.""" + meta = _build(_make_reader(ColumnInfo("a", "int64"))) + int(meta.schema_hash, 16) # raises ValueError if not hex + + +def test_schema_hash_changes_on_column_rename() -> None: + """Renaming a column changes the schema_hash.""" + m1 = _build(_make_reader(ColumnInfo("a", "int64"))) + m2 = _build(_make_reader(ColumnInfo("b", "int64"))) + assert m1.schema_hash != m2.schema_hash + + +def test_schema_hash_changes_on_dtype_change() -> None: + """Changing a column's dtype changes the schema_hash.""" + m1 = _build(_make_reader(ColumnInfo("a", "int64"))) + m2 = _build(_make_reader(ColumnInfo("a", "float64"))) + assert m1.schema_hash != m2.schema_hash + + +def test_schema_hash_stable_across_calls() -> None: + """Building metadata twice from the same reader yields the same schema_hash.""" + reader = _make_reader(ColumnInfo("x", "string"), ColumnInfo("y", "int64")) + m1 = _build(reader) + m2 = _build(reader) + assert m1.schema_hash == m2.schema_hash + + +def test_schema_hash_order_independent() -> None: + """schema_hash is unaffected by the order columns are declared in.""" + m1 = _build(_make_reader(ColumnInfo("a", "int64"), ColumnInfo("b", "string"))) + m2 = _build(_make_reader(ColumnInfo("b", "string"), ColumnInfo("a", "int64"))) + assert m1.schema_hash == m2.schema_hash + + +# --------------------------------------------------------------------------- +# source_hash +# --------------------------------------------------------------------------- + + +def test_explicit_source_hash_used_verbatim() -> None: + """An explicitly provided source_hash is stored as-is, not recomputed.""" + meta = _build(_make_reader(), source_hash="abc123") + assert meta.source_hash == "abc123" + + +def test_source_hash_computed_from_file(tmp_path: Path) -> None: + """Without an explicit hash, source_hash is the SHA-256 of the file contents.""" + content = b"col\n1\n2\n" + p = tmp_path / "ds.csv" + p.write_bytes(content) + expected = hashlib.sha256(content).hexdigest() + + meta = _build(_make_reader(), path=p) + assert meta.source_hash == expected + + +def test_source_hash_empty_for_nonexistent_path() -> None: + """source_hash is an empty string when the dataset path doesn't exist.""" + meta = _build(_make_reader(), path="/nonexistent/path.csv") + assert meta.source_hash == "" + + +# --------------------------------------------------------------------------- +# Each build produces a unique snapshot_id +# --------------------------------------------------------------------------- + + +def test_snapshot_ids_differ_across_builds() -> None: + """Repeated builds from the same reader each get a unique snapshot_id.""" + reader = _make_reader() + ids = {_build(reader).snapshot_id for _ in range(5)} + assert len(ids) == 5 diff --git a/tests/test_snapshot_stats.py b/tests/test_snapshot_stats.py new file mode 100644 index 0000000..419de91 --- /dev/null +++ b/tests/test_snapshot_stats.py @@ -0,0 +1,364 @@ +"""Tests for StatsComputer.""" + +from collections.abc import Iterator +from dataclasses import fields + +import pytest + +from dstrack.readers import Cell, ColumnInfo, TabularReader +from dstrack.snapshot import DatasetStats, StatsComputer + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _reader(cols: list[ColumnInfo], rows: list[list[Cell]]) -> TabularReader: + class _Stub: + def columns(self) -> list[ColumnInfo]: + return cols + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + if rows: + yield rows + + return _Stub() # type: ignore[return-value] + + +def _compute(cols: list[ColumnInfo], rows: list[list[Cell]]) -> DatasetStats: + return StatsComputer().compute(_reader(cols, rows)) + + +# --------------------------------------------------------------------------- +# num_rows +# --------------------------------------------------------------------------- + + +def test_num_rows_empty() -> None: + """num_rows is 0 for a dataset with no rows.""" + stats = _compute([ColumnInfo("x", "int64")], []) + assert stats.num_rows == 0 + + +def test_num_rows_counted() -> None: + """num_rows matches the number of rows yielded by the reader.""" + stats = _compute( + [ColumnInfo("x", "int64")], + [[1], [2], [3]], + ) + assert stats.num_rows == 3 + + +# --------------------------------------------------------------------------- +# duplicate_row_fraction +# --------------------------------------------------------------------------- + + +def test_no_duplicates() -> None: + """duplicate_row_fraction is 0.0 when all rows are distinct.""" + stats = _compute( + [ColumnInfo("x", "int64")], + [[1], [2], [3]], + ) + assert stats.duplicate_row_fraction == 0.0 + + +def test_all_identical_rows() -> None: + """duplicate_row_fraction counts all but the first occurrence as duplicates.""" + stats = _compute( + [ColumnInfo("x", "int64")], + [[5], [5], [5], [5]], + ) + # 3 out of 4 rows are duplicates + assert stats.duplicate_row_fraction == pytest.approx(3 / 4, abs=1e-9) + + +def test_partial_duplicates() -> None: + """duplicate_row_fraction reflects duplicates across multi-column rows.""" + stats = _compute( + [ColumnInfo("a", "int64"), ColumnInfo("b", "string")], + [[1, "x"], [1, "x"], [2, "y"]], + ) + assert stats.duplicate_row_fraction == pytest.approx(1 / 3, abs=1e-9) + + +def test_empty_dataset_duplicate_fraction() -> None: + """duplicate_row_fraction is 0.0 for a dataset with no rows.""" + stats = _compute([ColumnInfo("x", "int64")], []) + assert stats.duplicate_row_fraction == 0.0 + + +# --------------------------------------------------------------------------- +# high_null_columns +# --------------------------------------------------------------------------- + + +def test_high_null_columns_detected() -> None: + """A column with a high null fraction is listed in high_null_columns.""" + cols = [ColumnInfo("a", "int64"), ColumnInfo("b", "string")] + rows: list[list[Cell]] = [ + [1, None], + [2, None], + [3, None], + [4, "ok"], + ] + stats = _compute(cols, rows) + assert "b" in stats.high_null_columns + assert "a" not in stats.high_null_columns + + +def test_no_high_null_columns() -> None: + """high_null_columns is empty when no column has a high null fraction.""" + stats = _compute( + [ColumnInfo("x", "int64")], + [[1], [2]], + ) + assert stats.high_null_columns == [] + + +# --------------------------------------------------------------------------- +# constant_columns +# --------------------------------------------------------------------------- + + +def test_constant_numeric_column() -> None: + """A numeric column with a single repeated value is flagged as constant.""" + stats = _compute( + [ColumnInfo("x", "int64"), ColumnInfo("y", "int64")], + [[7, 1], [7, 2], [7, 3]], + ) + assert "x" in stats.constant_columns + assert "y" not in stats.constant_columns + + +def test_constant_string_column() -> None: + """A string column with a single repeated value is flagged as constant.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["hello"], ["hello"], ["hello"]], + ) + assert "s" in stats.constant_columns + + +def test_non_constant_string_column() -> None: + """A string column with varying values is not flagged as constant.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["a"], ["b"]], + ) + assert "s" not in stats.constant_columns + + +def test_all_null_column_is_constant() -> None: + """A column that is entirely null is treated as constant.""" + stats = _compute( + [ColumnInfo("n", "int64")], + [[None], [None]], + ) + assert "n" in stats.constant_columns + + +# --------------------------------------------------------------------------- +# Numeric column stats +# --------------------------------------------------------------------------- + + +def test_numeric_basic_stats() -> None: + """min/max/mean/null_count/null_fraction/num_unique are computed correctly.""" + stats = _compute( + [ColumnInfo("v", "int64")], + [[1], [2], [3], [4], [5]], + ) + s = stats.column_stats["v"] + assert s.min == 1.0 + assert s.max == 5.0 + assert s.mean == pytest.approx(3.0, abs=1e-9) + assert s.null_count == 0 + assert s.null_fraction == pytest.approx(0.0, abs=1e-9) + assert s.num_unique == 5 + + +def test_numeric_null_tracking() -> None: + """Null values in a numeric column are counted and reflected in null_fraction.""" + stats = _compute( + [ColumnInfo("v", "float64")], + [[1.0], [None], [3.0]], + ) + s = stats.column_stats["v"] + assert s.null_count == 1 + assert s.null_fraction == pytest.approx(1 / 3, abs=1e-9) + + +def test_numeric_std_single_value() -> None: + """std is 0.0 for a column with a single non-null value.""" + stats = _compute( + [ColumnInfo("v", "int64")], + [[42]], + ) + assert stats.column_stats["v"].std == pytest.approx(0.0, abs=1e-9) + + +def test_numeric_percentiles_present() -> None: + """All expected percentile fields are populated with plausible values.""" + stats = _compute( + [ColumnInfo("v", "int64")], + [[i] for i in range(1, 101)], + ) + p = stats.column_stats["v"].percentiles + assert {f.name for f in fields(p)} == {"p5", "p25", "p50", "p75", "p95", "p99"} + assert p.p50 == pytest.approx(50.0, abs=2) + assert p.p5 == pytest.approx(5.0, abs=2) + assert p.p25 == pytest.approx(25.0, abs=2) + assert p.p75 == pytest.approx(75.0, abs=2) + assert p.p95 == pytest.approx(95.0, abs=2) + assert p.p99 == pytest.approx(99.0, abs=2) + + +def test_numeric_histogram_bin_edge_count() -> None: + """Histogram bin_edges has one more entry than counts, and counts sum to num_rows.""" + rows = [[i] for i in range(1, 101)] + stats = _compute( + [ColumnInfo("v", "int64")], + rows, + ) + hist = stats.column_stats["v"].histogram + counts = hist.counts + edges = hist.bin_edges + assert len(edges) == len(counts) + 1 + assert sum(counts) == len(rows) + + +def test_numeric_all_null_returns_zeros() -> None: + """An all-null numeric column reports zeroed-out mean/std/num_unique.""" + stats = _compute( + [ColumnInfo("v", "int64")], + [[None], [None]], + ) + s = stats.column_stats["v"] + assert s.mean == 0.0 + assert s.std == 0.0 + assert s.num_unique == 0 + + +def test_bool_treated_as_numeric() -> None: + """A bool column is computed using numeric stats with 0/1 min/max.""" + stats = _compute( + [ColumnInfo("flag", "bool")], + [[True], [False], [True]], + ) + s = stats.column_stats["flag"] + assert s.mean == pytest.approx(2 / 3, abs=1e-9) + assert s.min == 0.0 + assert s.max == 1.0 + + +# --------------------------------------------------------------------------- +# String column stats +# --------------------------------------------------------------------------- + + +def test_string_basic_stats() -> None: + """null_count, num_unique, top_values, and top_values_coverage are correct.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["hello"], ["world"], ["hello"]], + ) + s = stats.column_stats["s"] + assert s.null_count == 0 + assert s.num_unique == 2 + assert s.top_values == {"hello": 2, "world": 1} + assert s.top_values_coverage == pytest.approx(1.0, abs=1e-9) + + +def test_string_char_length_stats() -> None: + """min/max/avg char length are computed across string values.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["a"], ["bb"], ["ccc"]], + ) + s = stats.column_stats["s"] + assert s.min_char_length == 1 + assert s.max_char_length == 3 + assert s.avg_char_length == pytest.approx(2.0, abs=1e-9) + + +def test_string_token_count_stats() -> None: + """min/max/avg token count are computed by splitting on whitespace.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["one"], ["two three"], ["four five six"]], + ) + s = stats.column_stats["s"] + assert s.min_token_count == 1 + assert s.max_token_count == 3 + assert s.avg_token_count == pytest.approx(2.0, abs=1e-9) + + +def test_string_null_tracking() -> None: + """Null values in a string column are counted and reflected in null_fraction.""" + stats = _compute( + [ColumnInfo("s", "string")], + [["a"], [None], [None]], + ) + s = stats.column_stats["s"] + assert s.null_count == 2 + assert s.null_fraction == pytest.approx(2 / 3, abs=1e-9) + + +def test_string_top_values_coverage_partial() -> None: + """top_values_coverage reflects the fraction covered by the top-N cap when unique values exceed it.""" + cols = [ColumnInfo("s", "string")] + rows: list[list[Cell]] = [[str(i)] for i in range(200)] + stats = _compute(cols, rows) + s = stats.column_stats["s"] + # 50 top values out of 200 = 0.25 coverage + assert s.top_values_coverage == pytest.approx(50 / 200, abs=1e-9) + + +# --------------------------------------------------------------------------- +# Datetime column stats +# --------------------------------------------------------------------------- + + +def test_datetime_min_max() -> None: + """min/max are the earliest/latest dates regardless of input order.""" + stats = _compute( + [ColumnInfo("dt", "datetime64")], + [["2024-01-01"], ["2024-06-15"], ["2023-12-31"]], + ) + s = stats.column_stats["dt"] + assert s.min == "2023-12-31" + assert s.max == "2024-06-15" + + +def test_datetime_range_days() -> None: + """range_days is the day span between the min and max dates.""" + stats = _compute( + [ColumnInfo("dt", "datetime64")], + [["2024-01-01"], ["2024-01-11"]], + ) + s = stats.column_stats["dt"] + assert s.range_days == pytest.approx(10.0, abs=1e-9) + + +def test_datetime_null_tracking() -> None: + """Null values in a datetime column are counted and reflected in null_fraction.""" + stats = _compute( + [ColumnInfo("dt", "datetime64")], + [["2024-01-01"], [None]], + ) + s = stats.column_stats["dt"] + assert s.null_count == 1 + assert s.null_fraction == pytest.approx(0.5, abs=1e-9) + + +def test_datetime_all_null() -> None: + """An all-null datetime column reports empty min/max and zero range_days.""" + stats = _compute( + [ColumnInfo("dt", "datetime64")], + [[None], [None]], + ) + s = stats.column_stats["dt"] + assert s.min == "" + assert s.max == "" + assert s.range_days == pytest.approx(0.0, abs=1e-9) From 38231d3ef216cc8ce6210221db533c334f95e6a0 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Sat, 11 Jul 2026 18:01:30 +0200 Subject: [PATCH 02/11] Creates a benchmark submodule that computes statistics on the speed at which a snapshot is created. --- pyproject.toml | 1 + src/dstrack/_benchmark/__init__.py | 54 +++++++ src/dstrack/_benchmark/_cli.py | 101 +++++++++++++ src/dstrack/_benchmark/_environment.py | 143 ++++++++++++++++++ src/dstrack/_benchmark/_profiling.py | 194 +++++++++++++++++++++++++ src/dstrack/_benchmark/_report.py | 152 +++++++++++++++++++ src/dstrack/_benchmark/_runner.py | 190 ++++++++++++++++++++++++ src/dstrack/_benchmark/_synthetic.py | 132 +++++++++++++++++ 8 files changed, 967 insertions(+) create mode 100644 src/dstrack/_benchmark/__init__.py create mode 100644 src/dstrack/_benchmark/_cli.py create mode 100644 src/dstrack/_benchmark/_environment.py create mode 100644 src/dstrack/_benchmark/_profiling.py create mode 100644 src/dstrack/_benchmark/_report.py create mode 100644 src/dstrack/_benchmark/_runner.py create mode 100644 src/dstrack/_benchmark/_synthetic.py diff --git a/pyproject.toml b/pyproject.toml index c47eb80..602d2b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ GitHub = "https://github.com/leoyala/dstrack" [project.scripts] dstrack = "dstrack._cli:app" +dstrack-benchmark = "dstrack._benchmark:app" [build-system] requires = ["uv_build>=0.11.26,<0.12.0"] diff --git a/src/dstrack/_benchmark/__init__.py b/src/dstrack/_benchmark/__init__.py new file mode 100644 index 0000000..5161550 --- /dev/null +++ b/src/dstrack/_benchmark/__init__.py @@ -0,0 +1,54 @@ +"""Benchmark for dstrack's snapshot creation pipeline. + +Generates a synthetic CSV in a temporary directory, then times how long it +takes to build snapshot metadata and statistics for it, alongside a summary of +the environment the benchmark ran in and a call tree of dstrack's own methods. + +Layers +------ +:mod:`._synthetic` + Describes and writes the synthetic dataset. + +:mod:`._runner` + Runs the snapshot pipeline against it and times each phase. + +:mod:`._profiling` + Reduces the ``cProfile`` run to a call tree of dstrack's own methods. + +:mod:`._environment` + Describes the machine the benchmark ran on. + +:mod:`._report`, :mod:`._cli` + Render runs to a console, and wire the whole thing to ``dstrack-benchmark``. +""" + +from ._cli import app +from ._environment import EnvironmentInfo, collect_environment_info +from ._profiling import CallGraph, CallNode, CallScope +from ._report import ConsoleReporter +from ._runner import ( + BenchmarkObserver, + BenchmarkResult, + BenchmarkRun, + BenchmarkRunner, + SilentObserver, +) +from ._synthetic import SyntheticColumn, SyntheticCsvSpec, write_synthetic_csv + +__all__ = [ + "BenchmarkObserver", + "BenchmarkResult", + "BenchmarkRun", + "BenchmarkRunner", + "CallGraph", + "CallNode", + "CallScope", + "ConsoleReporter", + "EnvironmentInfo", + "SilentObserver", + "SyntheticColumn", + "SyntheticCsvSpec", + "app", + "collect_environment_info", + "write_synthetic_csv", +] diff --git a/src/dstrack/_benchmark/_cli.py b/src/dstrack/_benchmark/_cli.py new file mode 100644 index 0000000..4c00f57 --- /dev/null +++ b/src/dstrack/_benchmark/_cli.py @@ -0,0 +1,101 @@ +"""Command-line entry point for the snapshot benchmark.""" + +import shutil +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Annotated + +import typer + +from dstrack._benchmark._report import ConsoleReporter +from dstrack._benchmark._runner import BenchmarkRunner +from dstrack._benchmark._synthetic import SyntheticCsvSpec + +app = typer.Typer( + no_args_is_help=True, + pretty_exceptions_show_locals=False, + help="Benchmark dstrack's snapshot creation pipeline against a synthetic CSV.", +) + +_DEFAULTS = SyntheticCsvSpec() + + +@contextmanager +def _workspace(reporter: ConsoleReporter, *, keep_csv: bool) -> Iterator[Path]: + """Yield a path for the synthetic CSV, discarding it unless ``keep_csv``.""" + tmp_dir = Path(tempfile.mkdtemp(prefix="dstrack-benchmark-")) + csv_path = tmp_dir / "synthetic_dataset.csv" + try: + yield csv_path + finally: + if keep_csv: + reporter.csv_kept(csv_path) + else: + shutil.rmtree(tmp_dir, ignore_errors=True) + + +@app.command(help="Generate a synthetic CSV and benchmark snapshot creation on it.") +def run( + rows: Annotated[ + int, typer.Option(help="Number of data rows to generate.") + ] = _DEFAULTS.num_rows, + numeric_cols: Annotated[ + int, typer.Option(help="Number of numeric (float) columns.") + ] = _DEFAULTS.num_numeric_cols, + string_cols: Annotated[ + int, typer.Option(help="Number of string columns.") + ] = _DEFAULTS.num_string_cols, + datetime_cols: Annotated[ + int, typer.Option(help="Number of datetime columns.") + ] = _DEFAULTS.num_datetime_cols, + bool_cols: Annotated[ + int, typer.Option(help="Number of boolean columns.") + ] = _DEFAULTS.num_bool_cols, + null_rate: Annotated[ + float, typer.Option(help="Fraction of non-id cells left null.") + ] = _DEFAULTS.null_rate, + seed: Annotated[ + int, typer.Option(help="Random seed for reproducible synthetic data.") + ] = _DEFAULTS.seed, + keep_csv: Annotated[ + bool, + typer.Option( + help="Keep the generated CSV instead of deleting it after the run." + ), + ] = False, + profile: Annotated[ + bool, + typer.Option( + help="Profile dstrack's own methods while building the snapshot. " + "Adds cProfile overhead to the reported metadata/stats timings." + ), + ] = True, + profile_limit: Annotated[ + int, + typer.Option( + help="Max number of child methods shown under each node of the " + "profile call tree, ranked by cumulative time." + ), + ] = 20, +) -> None: + spec = SyntheticCsvSpec( + num_rows=rows, + num_numeric_cols=numeric_cols, + num_string_cols=string_cols, + num_datetime_cols=datetime_cols, + num_bool_cols=bool_cols, + null_rate=null_rate, + seed=seed, + ) + reporter = ConsoleReporter() + runner = BenchmarkRunner(spec, observer=reporter, profile=profile) + + with _workspace(reporter, keep_csv=keep_csv) as csv_path: + benchmark_run = runner.run(csv_path) + reporter.report(benchmark_run, profile_limit=profile_limit) + + +if __name__ == "__main__": + app() diff --git a/src/dstrack/_benchmark/_environment.py b/src/dstrack/_benchmark/_environment.py new file mode 100644 index 0000000..02f84eb --- /dev/null +++ b/src/dstrack/_benchmark/_environment.py @@ -0,0 +1,143 @@ +"""Describes the machine and runtime a benchmark executed under. + +Hardware facts have no portable API, so each is read by a per-OS probe. Every +probe is best-effort: it returns ``None`` (or raises, which :func:`_attempt` +swallows) when the platform will not give up the answer, and the caller falls +back to something always available. +""" + +import os +import platform +import subprocess +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import TypeVar + +from dstrack import __version__ + +T = TypeVar("T") + +_PROBE_TIMEOUT_SECONDS = 5 + + +@dataclass(frozen=True) +class EnvironmentInfo: + """Snapshot of the machine and runtime a benchmark executed under.""" + + os_name: str + os_release: str + os_version: str + machine: str + cpu_model: str + cpu_count: int | None + total_memory_gb: float | None + python_implementation: str + python_version: str + dstrack_version: str + + +def _attempt(probe: Callable[[], T | None]) -> T | None: + """Run a probe, returning ``None`` if the platform refuses to answer.""" + try: + return probe() + except (OSError, subprocess.SubprocessError, ValueError): + return None + + +def _command_output(*args: str) -> str: + """Gets the stdout of a command, raising if it fails or times out.""" + result = subprocess.run( + args, + capture_output=True, + text=True, + timeout=_PROBE_TIMEOUT_SECONDS, + check=True, + ) + return result.stdout.strip() + + +def _first_matching_line(path: Path, prefix: str) -> str | None: + """First line of ``path`` starting with ``prefix`` (case-insensitive).""" + if not path.is_file(): + return None + for line in path.read_text().splitlines(): + if line.lower().startswith(prefix.lower()): + return line + return None + + +def _linux_cpu_model() -> str | None: + line = _first_matching_line(Path("/proc/cpuinfo"), "model name") + return line.split(":", 1)[1].strip() if line else None + + +def _darwin_cpu_model() -> str | None: + return _command_output("sysctl", "-n", "machdep.cpu.brand_string") + + +def _windows_cpu_model() -> str | None: + lines = [ + line.strip() + for line in _command_output("wmic", "cpu", "get", "name").splitlines() + if line.strip() + ] + return lines[1] if len(lines) >= 2 else None + + +def _linux_total_memory() -> int | None: + line = _first_matching_line(Path("/proc/meminfo"), "MemTotal:") + return int(line.split()[1]) * 1024 if line else None + + +def _darwin_total_memory() -> int | None: + return int(_command_output("sysctl", "-n", "hw.memsize")) + + +def _sysconf_total_memory() -> int | None: + """Page size times page count, on any POSIX platform exposing sysconf.""" + if not hasattr(os, "sysconf") or "SC_PAGE_SIZE" not in os.sysconf_names: + return None + return os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") + + +_CPU_MODEL_PROBES: dict[str, Callable[[], str | None]] = { + "Linux": _linux_cpu_model, + "Darwin": _darwin_cpu_model, + "Windows": _windows_cpu_model, +} + +_TOTAL_MEMORY_PROBES: dict[str, Callable[[], int | None]] = { + "Linux": _linux_total_memory, + "Darwin": _darwin_total_memory, +} + + +def _cpu_model() -> str: + """Best-effort human-readable CPU model string for the current machine.""" + probe = _CPU_MODEL_PROBES.get(platform.system()) + model = _attempt(probe) if probe is not None else None + return model or platform.processor() or platform.machine() or "unknown" + + +def _total_memory_bytes() -> int | None: + """Best-effort total physical memory in bytes, or None if it can't be read.""" + probe = _TOTAL_MEMORY_PROBES.get(platform.system(), _sysconf_total_memory) + return _attempt(probe) + + +def collect_environment_info() -> EnvironmentInfo: + """Collect OS, hardware, and runtime details describing the current machine.""" + total_memory = _total_memory_bytes() + return EnvironmentInfo( + os_name=platform.system(), + os_release=platform.release(), + os_version=platform.version(), + machine=platform.machine(), + cpu_model=_cpu_model(), + cpu_count=os.cpu_count(), + total_memory_gb=total_memory / (1024**3) if total_memory is not None else None, + python_implementation=platform.python_implementation(), + python_version=platform.python_version(), + dstrack_version=__version__, + ) diff --git a/src/dstrack/_benchmark/_profiling.py b/src/dstrack/_benchmark/_profiling.py new file mode 100644 index 0000000..15b900e --- /dev/null +++ b/src/dstrack/_benchmark/_profiling.py @@ -0,0 +1,194 @@ +"""Reduces a ``cProfile`` run to a call tree of dstrack's own methods. + +This module is pure model: it turns raw ``pstats`` frames into +:class:`CallNode` trees and knows nothing about how they are displayed. +Rendering lives in :mod:`dstrack._benchmark._report`. +""" + +import cProfile +import pstats +from collections.abc import Iterable +from dataclasses import dataclass, field +from pathlib import Path +from typing import TypeAlias + +# (filename, lineno, funcname) — the key pstats uses to identify a profiled frame. +FrameKey: TypeAlias = tuple[str, int, str] +# (primitive_calls, total_calls, total_time, cumulative_time) +_FrameTotals: TypeAlias = tuple[int, int, float, float] +# What pstats.Stats.stats maps each frame to: its totals plus its callers. +_RawStats: TypeAlias = dict[ + FrameKey, tuple[int, int, float, float, dict[FrameKey, _FrameTotals]] +] + + +@dataclass(frozen=True) +class CallScope: + """The source files whose profiled frames are worth reporting. + + A frame is in scope when it lives under ``root`` and under none of the + ``exclude`` paths, which keeps the tree focused on the code being measured + instead of the interpreter and the benchmark harness itself. + + Attributes: + root: Directory every reported frame must live under. + exclude: Sub-directories of ``root`` to leave out. + """ + + root: Path + exclude: tuple[Path, ...] = () + + def contains(self, filename: str) -> bool: + """True if frames from ``filename`` belong in the call tree.""" + try: + resolved = Path(filename).resolve() + except OSError: + return False + if not resolved.is_relative_to(self.root): + return False + return not any(resolved.is_relative_to(path) for path in self.exclude) + + def location(self, filename: str, lineno: int) -> str: + """Render a frame's source location relative to ``root``.""" + return f"{Path(filename).resolve().relative_to(self.root)}:{lineno}" + + +def dstrack_scope() -> CallScope: + """Scope covering dstrack's own code, minus this benchmark package.""" + benchmark_package = Path(__file__).resolve().parent + return CallScope(root=benchmark_package.parent, exclude=(benchmark_package,)) + + +@dataclass(frozen=True) +class CallNode: + """One method in the call tree, with the callees it was measured calling. + + Attributes: + funcname: Name of the profiled function. + location: Source location, relative to the scope root. + num_calls: Total number of calls, including recursive ones. + total_seconds: Time spent in the function itself. + cumulative_seconds: Time spent in the function and everything it called. + children: Callees, ordered by cumulative time, descending. + hidden_children: Callees omitted to respect the tree's child limit. + recursive: True if this node already appears higher up the same path, + in which case its callees are not expanded again. + """ + + funcname: str + location: str + num_calls: int + total_seconds: float + cumulative_seconds: float + children: tuple["CallNode", ...] = () + hidden_children: int = 0 + recursive: bool = False + + +@dataclass +class CallGraph: + """Caller/callee edges among the profiled frames that fall within a scope. + + ``pstats`` records, for every frame, the frames that called it. This + inverts those edges so a method that calls other in-scope methods (e.g. + ``StatsComputer.compute`` calling ``_build_dataset_stats``) becomes their + parent, then projects the result into trees rooted at the methods nothing + in scope called. + """ + + _scope: CallScope + _totals: dict[FrameKey, _FrameTotals] = field(default_factory=dict) + _children: dict[FrameKey, list[FrameKey]] = field(default_factory=dict) + _roots: list[FrameKey] = field(default_factory=list) + + @classmethod + def from_profile(cls, profiler: cProfile.Profile, scope: CallScope) -> "CallGraph": + """Build the in-scope call graph of a profiler that has finished running.""" + raw: _RawStats = pstats.Stats(profiler).stats # type: ignore[attr-defined] + graph = cls(_scope=scope) + frames = {key for key in raw if scope.contains(key[0])} + callees: set[FrameKey] = set() + + for key in frames: + primitive_calls, num_calls, total_time, cumulative_time, callers = raw[key] + graph._totals[key] = ( + primitive_calls, + num_calls, + total_time, + cumulative_time, + ) + for caller in callers: + if caller in frames: + graph._children.setdefault(caller, []).append(key) + callees.add(key) + + graph._roots = graph._by_cumulative(frames - callees) + return graph + + @property + def is_empty(self) -> bool: + """True if the profiler captured no in-scope frames at all.""" + return not self._totals + + def trees(self, *, max_children: int) -> list[CallNode]: + """Project the graph into call trees, hottest root first. + + Args: + max_children: Maximum callees expanded under each node, ranked by + cumulative time. The rest are counted in ``hidden_children``. + """ + return [ + self._expand(root, ancestors=frozenset({root}), max_children=max_children) + for root in self._roots + ] + + def _by_cumulative(self, keys: Iterable[FrameKey]) -> list[FrameKey]: + return sorted(keys, key=lambda key: self._totals[key][3], reverse=True) + + def _expand( + self, + key: FrameKey, + *, + ancestors: frozenset[FrameKey], + max_children: int, + ) -> CallNode: + """Build the node for ``key``, expanding its hottest callees. + + ``ancestors`` guards against (mutual) recursion: a callee already on the + current path is emitted as a leaf flagged ``recursive`` rather than + expanded again, which would otherwise never terminate. + """ + callees = self._by_cumulative(self._children.get(key, [])) + shown, hidden = callees[:max_children], callees[max_children:] + children = tuple( + self._node(child, recursive=True) + if child in ancestors + else self._expand( + child, + ancestors=ancestors | {child}, + max_children=max_children, + ) + for child in shown + ) + return self._node(key, children=children, hidden_children=len(hidden)) + + def _node( + self, + key: FrameKey, + *, + children: tuple[CallNode, ...] = (), + hidden_children: int = 0, + recursive: bool = False, + ) -> CallNode: + filename, lineno, funcname = key + _primitive_calls, num_calls, total_time, cumulative_time = self._totals[key] + return CallNode( + funcname=funcname, + location=self._scope.location(filename, lineno), + num_calls=num_calls, + total_seconds=total_time, + cumulative_seconds=cumulative_time, + children=children, + hidden_children=hidden_children, + recursive=recursive, + ) diff --git a/src/dstrack/_benchmark/_report.py b/src/dstrack/_benchmark/_report.py new file mode 100644 index 0000000..ada0e3d --- /dev/null +++ b/src/dstrack/_benchmark/_report.py @@ -0,0 +1,152 @@ +"""Renders benchmark runs to a Rich console. + +This is the only layer that knows what a benchmark looks like on screen. It +implements :class:`~dstrack._benchmark._runner.BenchmarkObserver`, so the same +object reports progress during a run and the tables after it. +""" + +from pathlib import Path + +from rich.console import Console +from rich.table import Table + +from dstrack._benchmark._environment import EnvironmentInfo +from dstrack._benchmark._profiling import CallGraph, CallNode +from dstrack._benchmark._runner import BenchmarkResult, BenchmarkRun +from dstrack.console import console as default_console + +_LAST_CONNECTOR = "└── " +_MID_CONNECTOR = "├── " +_LAST_GUIDE = " " +_MID_GUIDE = "│ " + + +class ConsoleReporter: + """Prints a benchmark's progress and results as Rich tables. + + Args: + console: Destination console. Defaults to dstrack's shared console; + inject another to capture the output in tests. + """ + + def __init__(self, console: Console | None = None) -> None: + self._console = console or default_console + + def generating_csv(self, path: Path, num_rows: int) -> None: + self._console.print( + f"Generating synthetic CSV at {path} ({num_rows:,} rows)..." + ) + + def building_metadata(self) -> None: + self._console.print("Building snapshot metadata...") + + def computing_stats(self) -> None: + self._console.print("Computing snapshot statistics...") + + def report(self, run: BenchmarkRun, *, profile_limit: int) -> None: + """Print the environment, the timings, and the profile call tree. + + Args: + run: The finished benchmark run. + profile_limit: Max callees shown under each node of the call tree. + """ + self._console.print(_environment_table(run.environment)) + self._console.print(_results_table(run.result)) + if run.call_graph is not None: + self._console.print(_call_tree_table(run.call_graph, profile_limit)) + + def csv_kept(self, path: Path) -> None: + """Report that the generated CSV was left on disk.""" + self._console.print(f"Synthetic CSV kept at {path}") + + +def _environment_table(env: EnvironmentInfo) -> Table: + table = Table(title="Environment") + table.add_column("Field") + table.add_column("Value") + table.add_row("OS", f"{env.os_name} {env.os_release}") + table.add_row("OS version", env.os_version) + table.add_row("Machine", env.machine) + table.add_row("CPU model", env.cpu_model) + table.add_row("CPU count", str(env.cpu_count) if env.cpu_count else "unknown") + table.add_row( + "Total memory", + f"{env.total_memory_gb:.2f} GB" + if env.total_memory_gb is not None + else "unknown", + ) + table.add_row("Python", f"{env.python_implementation} {env.python_version}") + table.add_row("dstrack version", env.dstrack_version) + return table + + +def _results_table(result: BenchmarkResult) -> Table: + table = Table(title="Snapshot benchmark") + table.add_column("Metric") + table.add_column("Value") + table.add_row("Rows generated", f"{result.num_rows:,}") + table.add_row("Columns", str(result.num_columns)) + table.add_row("CSV size", f"{result.csv_size_bytes / (1024**2):.2f} MB") + table.add_row("CSV generation time", f"{result.generation_seconds:.3f} s") + table.add_row("Metadata build time", f"{result.metadata_seconds:.3f} s") + table.add_row("Stats compute time", f"{result.stats_seconds:.3f} s") + table.add_row("Total snapshot time", f"{result.total_snapshot_seconds:.3f} s") + table.add_row("Throughput", f"{result.rows_per_second:,.0f} rows/s") + table.add_row("schema_hash", result.schema_hash) + return table + + +def _call_tree_table(graph: CallGraph, profile_limit: int) -> Table: + """Render dstrack's call tree, hottest root first. + + Guide characters (``├──``, ``└──``, ``│``) are baked into the "Method" cell + as plain text so the Calls/Total/Cumulative columns line up as a normal + table regardless of tree depth. + """ + table = Table(title="dstrack call tree (by cumulative time)") + table.add_column("Method", overflow="fold", ratio=1) + table.add_column("Calls", justify="right", no_wrap=True) + table.add_column("Total (s)", justify="right", no_wrap=True) + table.add_column("Cumulative (s)", justify="right", no_wrap=True) + + if graph.is_empty: + table.add_row("(no dstrack frames captured)", "-", "-", "-") + return table + + roots = graph.trees(max_children=profile_limit) + for i, root in enumerate(roots): + _add_node_rows(table, root, prefix="", is_last=i == len(roots) - 1) + return table + + +def _add_node_rows(table: Table, node: CallNode, *, prefix: str, is_last: bool) -> None: + """Add a row for ``node`` and, recursively, for the callees under it.""" + connector = _LAST_CONNECTOR if is_last else _MID_CONNECTOR + table.add_row( + f"{prefix}{connector}{_format_node(node)}", + str(node.num_calls), + f"{node.total_seconds:.4f}", + f"{node.cumulative_seconds:.4f}", + ) + + child_prefix = prefix + (_LAST_GUIDE if is_last else _MID_GUIDE) + for i, child in enumerate(node.children): + child_is_last = i == len(node.children) - 1 and not node.hidden_children + _add_node_rows(table, child, prefix=child_prefix, is_last=child_is_last) + + if node.hidden_children: + table.add_row( + f"{child_prefix}{_LAST_CONNECTOR}[dim]... {node.hidden_children} more " + "(raise --profile-limit to see them)[/dim]", + "-", + "-", + "-", + ) + + +def _format_node(node: CallNode) -> str: + """Render a node's function name and source location.""" + label = f"[bold]{node.funcname}[/bold] [dim]{node.location}[/dim]" + if node.recursive: + label += " [dim italic](recursive, not expanded)[/dim italic]" + return label diff --git a/src/dstrack/_benchmark/_runner.py b/src/dstrack/_benchmark/_runner.py new file mode 100644 index 0000000..d10e384 --- /dev/null +++ b/src/dstrack/_benchmark/_runner.py @@ -0,0 +1,190 @@ +"""Runs the snapshot pipeline against a synthetic CSV and times each phase. + +The runner performs no output of its own: progress is announced through a +:class:`BenchmarkObserver` and everything measured is returned as a +:class:`BenchmarkRun`, so the same run can be driven from a test with no +console attached. +""" + +import cProfile +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +from dstrack._benchmark._environment import EnvironmentInfo, collect_environment_info +from dstrack._benchmark._profiling import CallGraph, dstrack_scope +from dstrack._benchmark._synthetic import SyntheticCsvSpec, write_synthetic_csv +from dstrack.readers import CsvReader +from dstrack.snapshot import MetadataBuilder, StatsComputer + +_DATASET_NAME = "benchmark-dataset" +_CREATED_BY = "dstrack-benchmark" + + +@dataclass(frozen=True) +class BenchmarkResult: + """Timings and sizing for one synthetic-CSV snapshot benchmark run. + + Attributes: + num_rows: Rows written to the synthetic CSV. + num_rows_read: Rows the statistics pass actually read back. + num_columns: Columns the reader inferred, including ``id``. + csv_size_bytes: Size of the generated CSV on disk. + generation_seconds: Time to write the synthetic CSV. + metadata_seconds: Time to build snapshot metadata from the reader. + stats_seconds: Time for the statistics data pass. + schema_hash: Schema hash of the snapshot that was built. + """ + + num_rows: int + num_rows_read: int + num_columns: int + csv_size_bytes: int + generation_seconds: float + metadata_seconds: float + stats_seconds: float + schema_hash: str + + @property + def total_snapshot_seconds(self) -> float: + """Time to build the snapshot: metadata plus statistics.""" + return self.metadata_seconds + self.stats_seconds + + @property + def rows_per_second(self) -> float: + """Rows processed per second of snapshot time.""" + return self.num_rows_read / self.total_snapshot_seconds + + +@dataclass(frozen=True) +class BenchmarkRun: + """Everything one benchmark run produced, ready to be reported.""" + + result: BenchmarkResult + environment: EnvironmentInfo + call_graph: CallGraph | None + + +class BenchmarkObserver(Protocol): + """Receives a benchmark's progress as each phase starts.""" + + def generating_csv(self, path: Path, num_rows: int) -> None: + """Called before the synthetic CSV is written.""" + ... + + def building_metadata(self) -> None: + """Called before snapshot metadata is built.""" + ... + + def computing_stats(self) -> None: + """Called before the statistics data pass.""" + ... + + +class SilentObserver: + """Observer that reports nothing; the default for programmatic runs.""" + + def generating_csv(self, path: Path, num_rows: int) -> None: + return None + + def building_metadata(self) -> None: + return None + + def computing_stats(self) -> None: + return None + + +@dataclass +class _Elapsed: + """Wall-clock seconds a phase took, filled in when its block exits.""" + + seconds: float = 0.0 + + +class BenchmarkRunner: + """Generates a synthetic CSV and times a snapshot being built from it. + + Metadata and statistics are timed separately, and share one profiler so the + call tree covers the whole snapshot pipeline. CSV generation is never + profiled: it measures the harness, not dstrack. + + Args: + spec: The synthetic dataset to benchmark against. + observer: Notified as each phase starts. Defaults to reporting nothing. + profile: Whether to profile dstrack's own methods while building the + snapshot. Adds cProfile overhead to the reported timings. + """ + + def __init__( + self, + spec: SyntheticCsvSpec, + *, + observer: BenchmarkObserver | None = None, + profile: bool = True, + ) -> None: + self._spec = spec + self._observer = observer or SilentObserver() + self._profiler = cProfile.Profile() if profile else None + + def run(self, csv_path: Path) -> BenchmarkRun: + """Generate the dataset at ``csv_path`` and benchmark a snapshot of it.""" + self._observer.generating_csv(csv_path, self._spec.num_rows) + with self._timed() as generation: + write_synthetic_csv(csv_path, self._spec) + csv_size_bytes = csv_path.stat().st_size + + reader = CsvReader(csv_path) + + self._observer.building_metadata() + with self._timed(profiled=True) as metadata_phase: + metadata = MetadataBuilder().build( + reader, + dataset_name=_DATASET_NAME, + dataset_path=csv_path, + source_type="file", + created_by=_CREATED_BY, + ) + + self._observer.computing_stats() + with self._timed(profiled=True) as stats_phase: + stats = StatsComputer().compute(reader) + + result = BenchmarkResult( + num_rows=self._spec.num_rows, + num_rows_read=stats.num_rows, + num_columns=metadata.num_columns, + csv_size_bytes=csv_size_bytes, + generation_seconds=generation.seconds, + metadata_seconds=metadata_phase.seconds, + stats_seconds=stats_phase.seconds, + schema_hash=metadata.schema_hash, + ) + return BenchmarkRun( + result=result, + environment=collect_environment_info(), + call_graph=self._call_graph(), + ) + + @contextmanager + def _timed(self, *, profiled: bool = False) -> Iterator[_Elapsed]: + """Time the block, optionally recording it into the shared profiler.""" + profiler = self._profiler if profiled else None + elapsed = _Elapsed() + if profiler is not None: + profiler.enable() + start = time.perf_counter() + try: + yield elapsed + finally: + elapsed.seconds = time.perf_counter() - start + if profiler is not None: + profiler.disable() + + def _call_graph(self) -> CallGraph | None: + """Get a call graph of the snapshot phases, or None if profiling was disabled.""" + if self._profiler is None: + return None + return CallGraph.from_profile(self._profiler, dstrack_scope()) diff --git a/src/dstrack/_benchmark/_synthetic.py b/src/dstrack/_benchmark/_synthetic.py new file mode 100644 index 0000000..8a082f6 --- /dev/null +++ b/src/dstrack/_benchmark/_synthetic.py @@ -0,0 +1,132 @@ +"""Synthetic CSV generation for the benchmark. + +The dataset is described by a :class:`SyntheticCsvSpec` value object and +written by :func:`write_synthetic_csv`. Each generated dtype is one +:class:`SyntheticColumn` whose ``generate`` callable renders a single cell, so +supporting a new dtype means adding one cell factory and one entry to +``_COLUMN_KINDS`` rather than threading another count through every layer. +""" + +import csv +import random +from collections.abc import Callable +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import TypeAlias + +# Renders one cell as the text that lands in the CSV. Cells are always +# generated from the spec's seeded RNG, so a spec plus a seed is reproducible. +CellFactory: TypeAlias = Callable[[random.Random], str] + +_WORD_POOL: tuple[str, ...] = ( + "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", + "india", "juliet", "kilo", "lima", "mike", "november", "oscar", "papa", + "quebec", "romeo", "sierra", "tango", "uniform", "victor", "whiskey", +) # fmt: skip + +_DATETIME_START = datetime(2020, 1, 1, tzinfo=UTC) +_DATETIME_END = datetime(2026, 1, 1, tzinfo=UTC) + + +def _numeric_cell(rng: random.Random) -> str: + return f"{rng.gauss(100, 25):.4f}" + + +def _string_cell(rng: random.Random) -> str: + return " ".join(rng.choices(_WORD_POOL, k=rng.randint(1, 4))) + + +def _datetime_cell(rng: random.Random) -> str: + span = (_DATETIME_END - _DATETIME_START).total_seconds() + return (_DATETIME_START + timedelta(seconds=rng.uniform(0, span))).isoformat() + + +def _bool_cell(rng: random.Random) -> str: + return rng.choice(["True", "False"]) + + +@dataclass(frozen=True) +class SyntheticColumn: + """One generated column: its header name and how to render a cell. + + Attributes: + name: Column header written to the CSV. + generate: Renders a single non-null cell from the seeded RNG. + """ + + name: str + generate: CellFactory + + +# (name prefix, cell factory) per dtype dstrack's CSV reader infers, in the +# order the columns appear in the file. +_COLUMN_KINDS: tuple[tuple[str, CellFactory], ...] = ( + ("numeric", _numeric_cell), + ("string", _string_cell), + ("datetime", _datetime_cell), + ("bool", _bool_cell), +) + + +@dataclass(frozen=True) +class SyntheticCsvSpec: + """Describes the synthetic dataset a benchmark run should be measured on. + + Attributes: + num_rows: Number of data rows to generate. + num_numeric_cols: Number of ``float64`` columns. + num_string_cols: Number of ``string`` columns. + num_datetime_cols: Number of ``datetime64`` columns. + num_bool_cols: Number of ``bool`` columns. + null_rate: Fraction of non-id cells left empty (null). + seed: Seed for the random generator, for reproducible datasets. + """ + + num_rows: int = 200_000 + num_numeric_cols: int = 5 + num_string_cols: int = 5 + num_datetime_cols: int = 2 + num_bool_cols: int = 2 + null_rate: float = 0.02 + seed: int = 42 + + def columns(self) -> list[SyntheticColumn]: + """Return the nullable, randomly generated columns, in file order. + + The leading ``id`` column is not included: it is written from the row + index and is never null. + """ + counts = ( + self.num_numeric_cols, + self.num_string_cols, + self.num_datetime_cols, + self.num_bool_cols, + ) + return [ + SyntheticColumn(f"{prefix}_{i}", generate) + for (prefix, generate), count in zip(_COLUMN_KINDS, counts, strict=True) + for i in range(count) + ] + + +def write_synthetic_csv(path: Path, spec: SyntheticCsvSpec) -> None: + """Write the dataset described by ``spec`` to ``path``, overwriting it. + + Args: + path: Destination CSV path; overwritten if it already exists. + spec: The dataset to generate. + """ + rng = random.Random(spec.seed) + columns = spec.columns() + + with path.open("w", newline="", encoding="utf-8") as fh: + writer = csv.writer(fh) + writer.writerow(["id", *(column.name for column in columns)]) + for row_id in range(spec.num_rows): + row = [str(row_id)] + row.extend( + "" if rng.random() < spec.null_rate else column.generate(rng) + for column in columns + ) + writer.writerow(row) From ffa828cf89df28e9bc1f4ecc20cc87318c85ad90 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Tue, 14 Jul 2026 19:37:40 +0200 Subject: [PATCH 03/11] Adds feature to create local snapshots give a dataset. --- .../0003-local-snapshot-store-layout.md | 39 +- src/dstrack/_cli.py | 6 + src/dstrack/_track.py | 125 ++++++ src/dstrack/errors.py | 8 + src/dstrack/readers/__init__.py | 46 ++- src/dstrack/readers/_csv.py | 21 +- src/dstrack/readers/_protocol.py | 45 +++ src/dstrack/readers/_registry.py | 205 ++++++++++ src/dstrack/readers/_resolve.py | 138 +++++++ src/dstrack/snapshot/__init__.py | 3 + src/dstrack/snapshot/_builder.py | 111 ++++++ src/dstrack/snapshot/_metadata.py | 20 +- src/dstrack/store.py | 270 +++++++++++++ tests/test_reader_registry.py | 366 ++++++++++++++++++ tests/test_snapshot_metadata.py | 14 + tests/test_store.py | 99 +++++ 16 files changed, 1506 insertions(+), 10 deletions(-) create mode 100644 src/dstrack/_track.py create mode 100644 src/dstrack/readers/_registry.py create mode 100644 src/dstrack/readers/_resolve.py create mode 100644 src/dstrack/snapshot/_builder.py create mode 100644 src/dstrack/store.py create mode 100644 tests/test_reader_registry.py create mode 100644 tests/test_store.py diff --git a/docs/decisions/0003-local-snapshot-store-layout.md b/docs/decisions/0003-local-snapshot-store-layout.md index 9f710ba..cf1dcdb 100644 --- a/docs/decisions/0003-local-snapshot-store-layout.md +++ b/docs/decisions/0003-local-snapshot-store-layout.md @@ -146,13 +146,45 @@ Given `dstrack snapshot [--name NAME] [--root DIR] [--dataset-id ID]`: new dataset, `--name` is required, and a new `dataset_id` (UUID4) is minted. 4. If continuing an existing lineage, read `HEAD` and use it as `parent_snapshot_id`; otherwise `parent_snapshot_id` is `null`. -5. Read the source with the reader inferred from the file extension, compute +5. Read the source with the reader inferred from the file extension, or with the one + named by `--reader` (see [Readers are chosen per invocation, never + persisted](#readers-are-chosen-per-invocation-never-persisted) below), compute schema/stats/hashes per ADR-0001, and mint a new `snapshot_id` (UUID4). 6. Write `snapshots/.json`, append the corresponding line to `log.jsonl` (including `dataset_path`), then atomically replace `HEAD` (write-temp + rename), in that order, so a crash never leaves `HEAD` pointing at a snapshot that wasn't fully written. +### Readers are chosen per invocation, never persisted +A reader is resolved fresh on every `dstrack snapshot` call: inferred from the file's +extension, or named explicitly by `--reader`, which accepts either a registered reader +name (`--reader csv`) or a `package.module:ClassName` spec that `dstrack` imports +directly (`--reader mypackage.readers:ExcelReader`). + +**The resolved reader is deliberately not part of a snapshot.** No `reader` field is +written to `snapshots/*.json` or to `log.jsonl`, and nothing in the store is ever read +back to decide which reader to use. This is a security boundary, not an oversight: +a `package.module:ClassName` spec is arbitrary import-by-name, which is arbitrary code +execution. That is perfectly safe as an argument the invoking user typed themselves, and +it is exactly how gunicorn, uvicorn and celery accept application objects. It stops being +safe the moment the same string is read from a file, because `.dstrack/` is *committed to +git* (see [Location](#location)) and therefore travels with the repo. A reader spec stored +in a snapshot would mean that cloning a repository and running `dstrack snapshot` executes +an importable path chosen by whoever wrote that commit, turning a pull request into a code +execution vector on every reviewer's and CI runner's machine. + +The rule is therefore: **a reader spec may travel from the user to `dstrack`, never from +the store to `dstrack`.** Anything that would need a per-dataset default reader must be +solved another way, by an installed plugin claiming the extension via the +`dstrack.readers` entry-point group (code the user chose to `pip install`, on the same +footing as any other dependency), never by resurrecting an import path out of committed +history. + +The cost is real and accepted: a dataset whose extension is not claimed by any installed +reader needs `--reader` on every invocation, and cannot be made to "just work" for a +teammate by committing that fact. Publishing or vendoring the reader as a plugin package +is the supported answer. + ### What's committed vs. local-only | Path | Committed? | Why | @@ -249,3 +281,8 @@ this machine) return the same answer regardless of which machine ran them. dataset history; this is an implementation invariant, not just a convention. - `cache/index.db` can be deleted at any time with no data loss, only a slower next search until it's rebuilt. +- No reader information is ever written to the store, and no future field may reintroduce + it: because `.dstrack/` is committed, an importable path read back out of it would + execute code chosen by whoever authored the commit. Readers reach `dstrack` from the + invoking user (`--reader`) or from installed plugins, never from committed history. Any + proposal for a per-dataset default reader supersedes this ADR rather than extending it. diff --git a/src/dstrack/_cli.py b/src/dstrack/_cli.py index 1e77782..fe979fa 100644 --- a/src/dstrack/_cli.py +++ b/src/dstrack/_cli.py @@ -7,6 +7,7 @@ from dstrack import __version__, console from dstrack._store_template import STORE_TEMPLATE, materialize +from dstrack._track import track from dstrack.errors import StoreInitError from dstrack.utils import get_invocation_path @@ -84,6 +85,11 @@ def init( raise +app.command(help="Compute a snapshot of a dataset and store it in the local store.")( + track +) + + @app.command(help="Print package version.") def version() -> None: print(__version__) diff --git a/src/dstrack/_track.py b/src/dstrack/_track.py new file mode 100644 index 0000000..ed201f9 --- /dev/null +++ b/src/dstrack/_track.py @@ -0,0 +1,125 @@ +"""The ``dstrack track`` command: snapshot a dataset into the local store.""" + +import getpass +import os +from pathlib import Path +from typing import Annotated + +import typer + +from dstrack import console +from dstrack.errors import ( + DatasetNotFoundError, + StoreCorruptionError, + StoreNotFoundError, +) +from dstrack.paths import resolve_store_root +from dstrack.readers import resolve_reader +from dstrack.snapshot import SnapshotBuilder +from dstrack.store import write_snapshot + + +def _default_created_by() -> str: + """Best-effort identifier for the user creating the snapshot.""" + try: + return getpass.getuser() + except Exception: # pragma: no cover - platform dependent + return "unknown" + + +def track( + path: Annotated[ + Path, + typer.Argument( + exists=True, + dir_okay=False, + help="Path to the dataset file to snapshot.", + ), + ], + name: Annotated[ + str | None, + typer.Option( + "--name", + help="Human-readable dataset name. Defaults to the file's stem.", + ), + ] = None, + reader: Annotated[ + str | None, + typer.Option( + "--reader", + help="Reader to use: a registered name (e.g. 'csv'), or " + "'package.module:ClassName' for a reader that is not installed as a " + "plugin. Inferred from the file extension when omitted.", + ), + ] = None, + root: Annotated[ + Path | None, + typer.Option( + "--root", + help="Path root the stored dataset_path is made relative to. " + "Defaults to the store root (the directory containing .dstrack/).", + ), + ] = None, + dataset_id: Annotated[ + str | None, + typer.Option( + "--dataset-id", + help="Continue this dataset's lineage explicitly, instead of " + "matching by path. Needed after a file is renamed or moved.", + ), + ] = None, + created_by: Annotated[ + str | None, + typer.Option("--created-by", help="Override the recorded author."), + ] = None, +) -> None: + """Compute a snapshot of a dataset and store it in the local store.""" + # Locate the .dstrack/ store; without one there is nowhere to write. + try: + store_root = resolve_store_root() + except StoreNotFoundError as e: + console.error(str(e)) + raise typer.Exit(code=1) from e + + # Record the dataset location relative to the path root, so the store stays + # portable across machines and checkouts. + path_root = root if root is not None else store_root.parent + dataset_path = Path( + os.path.relpath(path.resolve(), Path(path_root).resolve()) + ).as_posix() + + # Pick the reader: explicit --reader wins, otherwise infer from the extension. + try: + tabular_reader = resolve_reader(path, reader=reader) + except (ValueError, TypeError) as e: + hint = "" if reader is not None else " Use --reader to name one explicitly." + console.error(f"{e}{hint}") + raise typer.Exit(code=1) from e + + # Read the file and compute the snapshot's statistics and metadata. The + # portable relative path is recorded; the real file on disk is what gets read + # and hashed. + console.info(f"Reading {path} and computing snapshot...") + snapshot = SnapshotBuilder().build( + tabular_reader, + dataset_name=name or path.stem, + dataset_path=dataset_path, + source=path, + source_type="file", # TODO: Add a way to manage other sources like folders + created_by=created_by or _default_created_by(), + ) + + # Persist the snapshot, attaching it to an existing lineage when one matches. + try: + result = write_snapshot(snapshot, store_root=store_root, dataset_id=dataset_id) + except (DatasetNotFoundError, StoreCorruptionError) as e: + console.error(str(e)) + raise typer.Exit(code=1) from e + + # Report what was written and whether it started or extended a lineage. + lineage = "new dataset" if result.is_new_dataset else "continued lineage" + console.success( + f"Snapshot {result.snapshot_id} written ({lineage}, " + f"dataset {result.dataset_id})." + ) + console.info(f"Stored at {result.snapshot_path}") diff --git a/src/dstrack/errors.py b/src/dstrack/errors.py index 3c64dfd..02d942c 100644 --- a/src/dstrack/errors.py +++ b/src/dstrack/errors.py @@ -4,3 +4,11 @@ class StoreInitError(Exception): class StoreNotFoundError(Exception): """Raised when no local store root can be resolved.""" + + +class StoreCorruptionError(Exception): + """Raised when a file in the local store cannot be read back as written.""" + + +class DatasetNotFoundError(Exception): + """Raised when a named dataset does not exist in the local store.""" diff --git a/src/dstrack/readers/__init__.py b/src/dstrack/readers/__init__.py index 6cd190a..0a293c7 100644 --- a/src/dstrack/readers/__init__.py +++ b/src/dstrack/readers/__init__.py @@ -6,7 +6,8 @@ Extending --------- -Implement [TabularReader][dstrack.readers.TabularReader] on any class to create a custom reader. +Implement [TabularReader][dstrack.readers.TabularReader] on any class to create a +custom reader. That alone is enough to build a snapshot from Python: Examples: ```python @@ -21,9 +22,48 @@ True ``` + +To make a reader reachable *by name* as well - from the CLI, or by extension +inference - it must also satisfy +[ReaderFactory][dstrack.readers._protocol.ReaderFactory] (a ``from_path`` classmethod), and +be registered. There are two ways in, for two different situations: + +- **Shipping a reader in a package** others install: declare an entry point in + the ``dstrack.readers`` group, and set ``EXTENSIONS`` on the class. It is then + picked up automatically, and ``dstrack track data.parquet`` just works with + nothing extra typed: + + ```toml + [project.entry-points."dstrack.readers"] + parquet = "dstrack_parquet:ParquetReader" + ``` + +- **A reader in your own project**, not installed as a plugin: call + [register_reader][dstrack.readers._registry.register_reader] from Python, or name it on + the command line as ``--reader "mypackage.readers:ExcelReader"``. """ from ._csv import CsvReader -from ._protocol import Cell, ColumnInfo, TabularReader +from ._protocol import Cell, ColumnInfo, ReaderFactory, TabularReader +from ._registry import ( + ENTRY_POINT_GROUP, + available_readers, + known_extensions, + register_reader, +) +from ._resolve import load_reader_class, resolve_reader, resolve_reader_class -__all__ = ["Cell", "ColumnInfo", "CsvReader", "TabularReader"] +__all__ = [ + "ENTRY_POINT_GROUP", + "Cell", + "ColumnInfo", + "CsvReader", + "ReaderFactory", + "TabularReader", + "available_readers", + "known_extensions", + "load_reader_class", + "register_reader", + "resolve_reader", + "resolve_reader_class", +] diff --git a/src/dstrack/readers/_csv.py b/src/dstrack/readers/_csv.py index b5212fb..3dc608f 100644 --- a/src/dstrack/readers/_csv.py +++ b/src/dstrack/readers/_csv.py @@ -131,7 +131,7 @@ def _coerce(raw: str | None, dtype: str) -> Cell: class CsvReader: """Reads a CSV file using the standard-library ``csv`` module. - Satisfies [TabularReader][dstrack.readers.TabularReader] without inheriting from + Satisfies [TabularReader][dstrack.readers._protocol.TabularReader] without inheriting from it. Column dtypes are inferred from the first ``sample_rows`` data rows; every subsequent call to [columns()][dstrack.readers.CsvReader.columns] returns the cached result. @@ -172,6 +172,25 @@ class CsvReader: (e.g. ``delimiter=";"``, ``quotechar="'"``). """ + EXTENSIONS = (".csv",) + + @classmethod + def from_path(cls, path: str | Path) -> "CsvReader": + """Build a reader for ``path`` with default options. + + Satisfies [ReaderFactory][dstrack.readers._protocol.ReaderFactory], which is how + the registry and ``--reader`` construct a reader they only know by name. + Options other than the path are not reachable this way; construct the + reader directly to set them. + + Args: + path: Path to the CSV file. + + Returns: + A `CsvReader` bound to `path`. + """ + return cls(path) + def __init__( self, path: str | Path, diff --git a/src/dstrack/readers/_protocol.py b/src/dstrack/readers/_protocol.py index f07cfe9..2c872af 100644 --- a/src/dstrack/readers/_protocol.py +++ b/src/dstrack/readers/_protocol.py @@ -1,5 +1,6 @@ from collections.abc import Iterator from dataclasses import dataclass +from pathlib import Path from typing import Protocol, TypeAlias, runtime_checkable # A single cell value after dtype coercion; None represents a missing value. @@ -65,3 +66,47 @@ def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: A list of rows, each row being a list of [Cell][dstrack.readers.Cell] values. """ ... + + +@runtime_checkable +class ReaderFactory(Protocol): + """Construction contract for readers reached *by name* rather than by instance. + + [TabularReader][dstrack.readers.TabularReader] describes how a reader is + *read*, and says nothing about how one is *built*: code that already holds an + instance never needs to know. But the registry and the + ``"package.module:ClassName"`` spec only ever yield a class, so they need a + uniform way to turn that class into an instance given a source path. + ``from_path`` is that way, and it is checked against this protocol before the + class is ever called. + + This is deliberately a second, separate protocol: a reader used only from + Python (constructed by the caller, handed straight to ``SnapshotBuilder``) + still needs nothing beyond ``TabularReader``. Only readers that are + registered, or named on the command line, must also satisfy this one. + + Examples: + ```python + >>> from pathlib import Path + >>> from dstrack.readers import CsvReader, ReaderFactory + >>> isinstance(CsvReader, ReaderFactory) # the class, not an instance + True + + ``` + """ + + def from_path(self, path: Path) -> TabularReader: + """Build a reader for ``path`` using default options. + + Implemented as a ``classmethod`` on the reader class; the protocol is + therefore checked against the class object itself + (``isinstance(MyReader, ReaderFactory)``). + + Args: + path: Path to the dataset source the reader will read. + + Returns: + An instance satisfying + [TabularReader][dstrack.readers.TabularReader]. + """ + ... diff --git a/src/dstrack/readers/_registry.py b/src/dstrack/readers/_registry.py new file mode 100644 index 0000000..beb37eb --- /dev/null +++ b/src/dstrack/readers/_registry.py @@ -0,0 +1,205 @@ +"""Registry of reader classes, keyed by short name and by file extension. + +Registration is the single extension point for readers. Three kinds of user +reach it by three different routes, and none of them should have to do the other +two's work: + +- **From Python**, a caller who already has a reader instance hands it straight + to `SnapshotBuilder`; the registry is never consulted. +- **From an installed plugin package**, a reader self-registers through the + ``dstrack.readers`` entry-point group, so ``dstrack track data.parquet`` infers + it from the extension with nothing to type. +- **From an ad-hoc class** in the user's own project, which is not installed as a + plugin and so is named explicitly as ``"package.module:ClassName"`` (see + `_resolve`). + +A registered class must satisfy both +[TabularReader][dstrack.readers.TabularReader] (how it reads) and +[ReaderFactory][dstrack.readers.ReaderFactory] (how it is built from a path); +`check_reader_class` enforces that *before* the class is ever instantiated. +""" + +import logging +from collections.abc import Sequence +from importlib.metadata import entry_points +from pathlib import Path +from typing import Final + +from ._csv import CsvReader +from ._protocol import ReaderFactory, TabularReader + +_log = logging.getLogger(__name__) + +# Entry-point group an installed package declares to add a reader, e.g. +# [project.entry-points."dstrack.readers"] +# parquet = "dstrack_parquet:ParquetReader" +ENTRY_POINT_GROUP: Final[str] = "dstrack.readers" + +# Short name (as typed for --reader) -> reader class. +_BY_NAME: dict[str, type[TabularReader]] = {} +# Lowercase file extension, leading dot included -> reader class. +_BY_EXTENSION: dict[str, type[TabularReader]] = {} + +_plugins_loaded = False + + +def check_reader_class(obj: object, *, origin: str) -> type[TabularReader]: + """Validate that ``obj`` is a class usable as a reader, without calling it. + + Both protocols are checked here, on the class, so that a class which merely + *looks* like a reader is rejected before its constructor runs on a user path. + + Args: + obj: The candidate, typically freshly imported or freshly registered. + origin: How the caller referred to it (a spec, a name, an entry point), + used to make the error message actionable. + + Returns: + ``obj`` itself, narrowed to a reader class. + + Raises: + TypeError: If ``obj`` is not a class, or does not satisfy + [TabularReader][dstrack.readers.TabularReader] and + [ReaderFactory][dstrack.readers.ReaderFactory]. + """ + if not isinstance(obj, type): + raise TypeError( + f"{origin} resolved to {obj!r}, which is not a class. " + "A reader must be a class, not an instance or a function." + ) + qualname = f"{obj.__module__}.{obj.__qualname__}" + if not issubclass(obj, TabularReader): + raise TypeError( + f"{qualname} (from {origin}) does not satisfy the TabularReader " + "protocol: it must define columns() and iter_batches()." + ) + # Checked through an object-typed name: mypy assumes any class object + # already satisfies ReaderFactory, so narrowing from `object` is what keeps + # this branch live for the type checker as well as at runtime. + candidate: object = obj + if not isinstance(candidate, ReaderFactory): + raise TypeError( + f"{qualname} (from {origin}) does not satisfy the ReaderFactory " + "protocol: it must define a from_path(path) classmethod, which is " + "how dstrack builds a reader it knows only by name." + ) + return obj + + +def build_reader(reader_cls: type[TabularReader], path: Path) -> TabularReader: + """Instantiate a validated reader class for ``path``. + + Args: + reader_cls: A class already validated through `check_reader_class`. + path: Path to the dataset source. + + Returns: + The reader instance. + + Raises: + TypeError: If ``reader_cls`` was never validated and lacks ``from_path``. + """ + factory: object = reader_cls + if not isinstance(factory, ReaderFactory): + raise TypeError(f"{reader_cls!r} is not an instance of {ReaderFactory}") + return factory.from_path(path) + + +def register_reader( + reader_cls: type[TabularReader], + *, + name: str, + extensions: Sequence[str] | None = None, +) -> None: + """Register a reader under a short name and the extensions it handles. + + Args: + reader_cls: The reader class. Must satisfy both `TabularReader` and + `ReaderFactory`. + name: Short name, as typed for ``--reader`` (e.g. ``"csv"``). + extensions: Extensions this reader claims, leading dot included. When + omitted, the class's ``EXTENSIONS`` attribute is used, so a plugin + can declare its extensions once on the class and register through a + bare entry point. + + Raises: + TypeError: If ``reader_cls`` does not satisfy both reader protocols. + ValueError: If ``name`` or any extension is already taken. Registration + never silently displaces an existing reader: two packages fighting + over ``.parquet`` is a conflict the user has to see, not one to + resolve by import order. + """ + check_reader_class(reader_cls, origin=f"reader {name!r}") + + if extensions is None: + extensions = getattr(reader_cls, "EXTENSIONS", ()) + + if name in _BY_NAME: + raise ValueError( + f"Reader name {name!r} is already registered to " + f"{_BY_NAME[name].__module__}.{_BY_NAME[name].__qualname__}." + ) + + claimed = [ext.lower() for ext in extensions] + for ext in claimed: + if ext in _BY_EXTENSION: + owner = _BY_EXTENSION[ext] + raise ValueError( + f"Extension {ext!r} is already registered to " + f"{owner.__module__}.{owner.__qualname__}." + ) + + _BY_NAME[name] = reader_cls + for ext in claimed: + _BY_EXTENSION[ext] = reader_cls + + +def _load_plugins() -> None: + """Register readers advertised by installed packages, once per process. + + A plugin that fails to import or conflicts with an existing registration is + logged and skipped: one broken third-party package must not take down + `dstrack track` for a dataset that does not even use it. + """ + global _plugins_loaded + if _plugins_loaded: + return + # Set before loading: a plugin that raises must not be retried on every + # subsequent lookup. + _plugins_loaded = True + + for entry_point in entry_points(group=ENTRY_POINT_GROUP): + try: + register_reader(entry_point.load(), name=entry_point.name) + except Exception as e: + _log.warning( + f"Ignoring reader plugin {entry_point.name!r} " + f"({entry_point.value}): {e}" + ) + + +def reader_class_for_name(name: str) -> type[TabularReader] | None: + """Look up a reader by its short name, or ``None`` if nothing claims it.""" + _load_plugins() + return _BY_NAME.get(name) + + +def reader_class_for_extension(extension: str) -> type[TabularReader] | None: + """Look up a reader by file extension, or ``None`` if nothing claims it.""" + _load_plugins() + return _BY_EXTENSION.get(extension.lower()) + + +def available_readers() -> dict[str, type[TabularReader]]: + """Return all registered readers by short name, built-in and plugin alike.""" + _load_plugins() + return dict(_BY_NAME) + + +def known_extensions() -> dict[str, type[TabularReader]]: + """Return all registered extensions and the reader class that claims each.""" + _load_plugins() + return dict(_BY_EXTENSION) + + +register_reader(CsvReader, name="csv") diff --git a/src/dstrack/readers/_resolve.py b/src/dstrack/readers/_resolve.py new file mode 100644 index 0000000..9ca8286 --- /dev/null +++ b/src/dstrack/readers/_resolve.py @@ -0,0 +1,138 @@ +"""Selection and loading of a :class:`TabularReader` for a given source. + +A reader is chosen one of three ways, in order of how much the user has to type: + +1. Implicitly, from the source file's extension (``data.csv`` -> `CsvReader`). + Plugin readers registered through the ``dstrack.readers`` entry-point group + participate here too, so an installed reader needs nothing on the command line. +2. By short name (``"csv"``), for when the extension is absent or misleading. +3. By ``"package.module:ClassName"`` spec, which imports the class directly. This + is the escape hatch for a reader that lives in the user's own project and is + not installed as a plugin. + +Security: a spec is arbitrary import-by-name, i.e. code execution. It is only +ever accepted from the invoking user (a ``--reader`` argument they typed), and +is never persisted into a snapshot nor read back out of one. See ADR-0003. +""" + +import importlib +from pathlib import Path + +from ._protocol import TabularReader +from ._registry import ( + available_readers, + build_reader, + check_reader_class, + known_extensions, + reader_class_for_extension, + reader_class_for_name, +) + + +def load_reader_class(spec: str) -> type[TabularReader]: + """Import a reader class from a ``"package.module:ClassName"`` spec. + + The class is validated against both reader protocols before being returned, + so a mistyped spec that happens to name some other importable object fails + here rather than part-way through a snapshot. + + Args: + spec: A ``":"`` string. The module part is imported and + the class part is looked up on it. + + Returns: + The referenced class (not an instance). + + Raises: + ValueError: If ``spec`` is not of the form ``"module:ClassName"``, the + module cannot be imported, or the class is not found on it. + TypeError: If the referenced object is not a usable reader class. + """ + module_name, sep, class_name = spec.partition(":") + if not sep or not module_name or not class_name: + raise ValueError( + f"Invalid reader spec {spec!r}. " + "Expected the form 'package.module:ClassName'." + ) + try: + module = importlib.import_module(module_name) + except ImportError as e: + raise ValueError( + f"Could not import module {module_name!r} from reader spec {spec!r}." + ) from e + try: + obj = getattr(module, class_name) + except AttributeError as e: + raise ValueError( + f"Module {module_name!r} has no attribute {class_name!r} " + f"(from reader spec {spec!r})." + ) from e + + return check_reader_class(obj, origin=f"reader spec {spec!r}") + + +def resolve_reader_class( + path: str | Path, *, reader: str | None = None +) -> type[TabularReader]: + """Choose the reader class for ``path``, without instantiating it. + + Args: + path: Path to the dataset source. + reader: Either a registered short name (``"csv"``) or a + ``"package.module:ClassName"`` spec, told apart by the ``":"``. + When omitted, the reader is inferred from ``path``'s extension. + + Returns: + A validated reader class. + + Raises: + ValueError: If ``reader`` names nothing known, a spec is malformed, or + no reader is registered for ``path``'s extension. + TypeError: If the resolved class is not a usable reader. + """ + if reader is not None: + if ":" in reader: + return load_reader_class(reader) + reader_cls = reader_class_for_name(reader) + if reader_cls is None: + known = ", ".join(sorted(available_readers())) or "(none)" + raise ValueError( + f"Unknown reader {reader!r}. Available readers: {known}. " + "Pass 'package.module:ClassName' to use a reader that is not " + "installed as a plugin." + ) + return reader_cls + + path = Path(path) + extension = path.suffix.lower() + reader_cls = reader_class_for_extension(extension) + if reader_cls is None: + known = ", ".join(sorted(known_extensions())) or "(none)" + raise ValueError( + f"No reader registered for extension {extension!r} (path: {path}). " + f"Known extensions: {known}. Name a reader explicitly, either a " + "registered one or 'package.module:ClassName'." + ) + return reader_cls + + +def resolve_reader(path: str | Path, *, reader: str | None = None) -> TabularReader: + """Build a reader for ``path``, explicitly named or inferred from its extension. + + Args: + path: Path to the dataset source the reader will read. + reader: Optional short name or ``"package.module:ClassName"`` spec. When + omitted, the reader is inferred from ``path``'s file extension. + + Returns: + An instantiated reader satisfying + [TabularReader][dstrack.readers.TabularReader]. + + Raises: + ValueError: If ``reader`` names nothing known, a spec is malformed, or + no reader is registered for ``path``'s extension. + TypeError: If the resolved class is not a usable reader. + """ + path = Path(path) + reader_cls = resolve_reader_class(path, reader=reader) + return build_reader(reader_cls, path) diff --git a/src/dstrack/snapshot/__init__.py b/src/dstrack/snapshot/__init__.py index 9a2e9f6..3253082 100644 --- a/src/dstrack/snapshot/__init__.py +++ b/src/dstrack/snapshot/__init__.py @@ -13,6 +13,7 @@ :class:`SnapshotMetadata`, :class:`DatasetStats` """ +from ._builder import SnapshotBuilder, build_snapshot_dict from ._metadata import MetadataBuilder, SnapshotMetadata from ._stats import ( DatasetStats, @@ -33,7 +34,9 @@ "NumericColumnStats", "OtherColumnStats", "PercentileStats", + "SnapshotBuilder", "SnapshotMetadata", "StatsComputer", "StringColumnStats", + "build_snapshot_dict", ] diff --git a/src/dstrack/snapshot/_builder.py b/src/dstrack/snapshot/_builder.py new file mode 100644 index 0000000..7d5e562 --- /dev/null +++ b/src/dstrack/snapshot/_builder.py @@ -0,0 +1,111 @@ +import dataclasses +from pathlib import Path +from typing import Any + +from dstrack.readers import TabularReader +from dstrack.snapshot._metadata import MetadataBuilder, SnapshotMetadata +from dstrack.snapshot._stats import DatasetStats, StatsComputer + + +class SnapshotBuilder: + """Builds a complete dataset snapshot from a single reader. + + Combines [MetadataBuilder][dstrack.snapshot._metadata.MetadataBuilder] + (identity and schema) and + [StatsComputer][dstrack.snapshot._stats.StatsComputer] (a data pass over + the rows) into one JSON-ready snapshot dict. Import it to build snapshots + from Python without going through the CLI: + + Examples: + >>> from dstrack.readers import CsvReader + >>> from dstrack.snapshot import SnapshotBuilder + >>> reader = CsvReader("data.csv") # doctest: +SKIP + >>> snapshot = SnapshotBuilder().build( # doctest: +SKIP + ... reader, + ... dataset_name="customers", + ... dataset_path="data.csv", + ... created_by="alice", + ... ) + + Args: + metadata_builder: Builder used for identity and schema fields. + Defaults to a fresh + [MetadataBuilder][dstrack.snapshot._metadata.MetadataBuilder]. + stats_computer: Computer used for the per-column and dataset-level + statistics. Defaults to a fresh + [StatsComputer][dstrack.snapshot._stats.StatsComputer]. + """ + + def __init__( + self, + *, + metadata_builder: MetadataBuilder | None = None, + stats_computer: StatsComputer | None = None, + ) -> None: + self._metadata_builder = metadata_builder or MetadataBuilder() + self._stats_computer = stats_computer or StatsComputer() + + def build( + self, + reader: TabularReader, + *, + dataset_name: str, + dataset_path: str | Path, + created_by: str, + source_type: str = "file", + source: str | Path | None = None, + source_hash: str | None = None, + ) -> dict[str, Any]: + """Build a JSON-serializable snapshot from a reader. + + The reader is consumed once for schema inference and once for the + statistics data pass. + + Args: + reader: Any [TabularReader][dstrack.readers._protocol.TabularReader]. + dataset_name: Human-readable dataset name stored in the snapshot. + dataset_path: Source path or URI recorded in the snapshot. Never + opened; pass ``source`` when the data lives elsewhere. + created_by: User or process identifier. + source_type: Origin kind (``"file"``, ``"directory"``, etc.). + source: Location the data actually lives at, used only to compute + ``source_hash``. Defaults to ``dataset_path``. + source_hash: Pre-computed source hash. When ``None`` and the + source resolves to a regular file, a SHA-256 of the file bytes + is computed automatically by the metadata builder. + + Returns: + A ``dict`` merging every metadata and statistics field, ready to be + serialized with [json.dumps][json.dumps]. + """ + metadata = self._metadata_builder.build( + reader, + dataset_name=dataset_name, + dataset_path=dataset_path, + source_type=source_type, + created_by=created_by, + source=source, + source_hash=source_hash, + ) + stats = self._stats_computer.compute(reader) + return build_snapshot_dict(metadata, stats) + + +def build_snapshot_dict( + metadata: SnapshotMetadata, stats: DatasetStats +) -> dict[str, Any]: + """Merge metadata and statistics into one JSON-serializable snapshot dict. + + The two inputs contribute disjoint sets of keys, so the merge never drops + a field. + + Args: + metadata: Identity and schema fields from + [MetadataBuilder][dstrack.snapshot._metadata.MetadataBuilder]. + stats: Volume, per-column, and quality fields from + [StatsComputer][dstrack.snapshot._stats.StatsComputer]. + + Returns: + A ``dict`` combining both, with nested dataclasses expanded to dicts. + """ + return {**dataclasses.asdict(metadata), **dataclasses.asdict(stats)} diff --git a/src/dstrack/snapshot/_metadata.py b/src/dstrack/snapshot/_metadata.py index 8e56414..e9c349a 100644 --- a/src/dstrack/snapshot/_metadata.py +++ b/src/dstrack/snapshot/_metadata.py @@ -50,19 +50,29 @@ def build( dataset_path: str | Path, source_type: str, created_by: str, + source: str | Path | None = None, source_hash: str | None = None, ) -> SnapshotMetadata: """Build metadata for a snapshot. + ``dataset_path`` is what gets *recorded*; ``source`` is what gets + *read*. They differ whenever the recorded path is relative to a path + root that is not the current working directory, as it is for snapshots + written by the CLI. + Args: reader: Any TabularReader; only ``columns()`` is called. dataset_name: Human-readable dataset name stored in the snapshot. - dataset_path: Source path or URI at snapshot time. + dataset_path: Source path or URI at snapshot time, recorded in the + snapshot verbatim. Never opened. source_type: Origin kind (``"file"``, ``"directory"``, etc.). created_by: User or process identifier. - source_hash: Pre-computed source hash. When ``None`` and - ``dataset_path`` resolves to a regular file, a SHA-256 of - the file bytes is computed automatically. + source: Location the data actually lives at, used only to compute + ``source_hash``. Defaults to ``dataset_path``, which is + correct when the recorded path is one the process can open. + source_hash: Pre-computed source hash. When ``None`` and the + source resolves to a regular file, a SHA-256 of the file bytes + is computed automatically. Returns: A populated :class:`SnapshotMetadata` instance. @@ -71,7 +81,7 @@ def build( path_str = str(dataset_path) if source_hash is None: - p = Path(dataset_path) + p = Path(source if source is not None else dataset_path) source_hash = _hash_file(p) if p.is_file() else "" return SnapshotMetadata( diff --git a/src/dstrack/store.py b/src/dstrack/store.py new file mode 100644 index 0000000..a36bb95 --- /dev/null +++ b/src/dstrack/store.py @@ -0,0 +1,270 @@ +"""Persistence of snapshots into the local store. + +A snapshot is written under ``datasets//`` as its full JSON +payload, a one-line append to ``log.jsonl``, and an updated ``HEAD``. The +three are written in that order so a crash never leaves ``HEAD`` +pointing at a snapshot that was not fully written. Re-tracking a source whose +recorded path matches an existing dataset's latest snapshot continues that +dataset's lineage rather than creating a new one. +""" + +import json +import os +import tempfile +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Final + +from dstrack.errors import DatasetNotFoundError, StoreCorruptionError + +# Lightweight identity fields copied from a snapshot into its `log.jsonl` line. +# A subset of ADR-0001's fields: only those the current builders produce. +_LOG_FIELDS: Final[tuple[str, ...]] = ( + "snapshot_id", + "parent_snapshot_id", + "created_at", + "created_by", + "dataset_name", + "dataset_path", + "num_rows", + "num_columns", +) + + +@dataclass(frozen=True) +class SnapshotWriteResult: + """Outcome of writing one snapshot into the store. + + Attributes: + dataset_id: Dataset whose lineage the snapshot joined, whether it was + passed in, matched by path, or minted by this write. + snapshot_id: Identifier of the snapshot that was written, copied from + the snapshot payload. + snapshot_path: Path of the written ``snapshots/.json``. + parent_snapshot_id: Snapshot this one succeeds, i.e. the dataset's + ``HEAD`` before this write, or ``None`` for a dataset's first + snapshot. + is_new_dataset: Whether this write minted a new ``dataset_id``, as + opposed to continuing an existing dataset's lineage. + """ + + dataset_id: str + snapshot_id: str + snapshot_path: Path + parent_snapshot_id: str | None + is_new_dataset: bool + + +def write_snapshot( + snapshot: dict[str, Any], + *, + store_root: Path, + dataset_id: str | None = None, +) -> SnapshotWriteResult: + """Write ``snapshot`` into the store and update the dataset's history. + + Resolves which dataset the snapshot belongs to, then writes the full + payload, appends the lightweight ``log.jsonl`` line, and moves ``HEAD`` + onto the new snapshot. + + Args: + snapshot: A snapshot dict as produced by + [SnapshotBuilder][dstrack.snapshot.SnapshotBuilder]. Must contain + ``snapshot_id`` and ``dataset_path``; ``parent_snapshot_id`` is + filled in here and overwritten if already present. + store_root: Path to the ``.dstrack/`` directory, e.g. from + [resolve_store_root][dstrack.paths.resolve_store_root]. + dataset_id: Continue this dataset's lineage explicitly. The dataset must + already exist. When ``None``, the dataset is matched by + ``dataset_path`` against each existing dataset's latest snapshot; a + new ``dataset_id`` is minted if none matches. + + Returns: + A [SnapshotWriteResult][dstrack.store.SnapshotWriteResult] describing + where the snapshot landed and which lineage it joined. + + Raises: + KeyError: If ``snapshot`` lacks ``snapshot_id`` or ``dataset_path``. + DatasetNotFoundError: If ``dataset_id`` is given but names no dataset in + the store. + StoreCorruptionError: If an existing dataset's ``log.jsonl`` ends in a + line that is not valid JSON. + OSError: If the store cannot be written to. + """ + datasets_dir = store_root / "datasets" + dataset_path = snapshot["dataset_path"] + snapshot_id = snapshot["snapshot_id"] + + if dataset_id is not None: + # An explicit id continues a lineage; it must name a dataset that is + # already there, or a typo would mint a nameless one behind the user's + # back and report it as a continuation. + _check_dataset_exists(datasets_dir, dataset_id) + is_new_dataset = False + else: + dataset_id = _match_dataset_by_path(datasets_dir, dataset_path) + is_new_dataset = dataset_id is None + if dataset_id is None: + dataset_id = str(uuid.uuid4()) + + dataset_dir = datasets_dir / dataset_id + snapshots_dir = dataset_dir / "snapshots" + snapshots_dir.mkdir(parents=True, exist_ok=True) + + # HEAD is the dataset's latest snapshot, so it is this snapshot's parent. + # Absent for a dataset's first snapshot, whether it was just minted here or + # named by an explicit `dataset_id` that has no snapshots yet. + parent_snapshot_id = _read_head(dataset_dir) + payload = {**snapshot, "parent_snapshot_id": parent_snapshot_id} + + snapshot_path = snapshots_dir / f"{snapshot_id}.json" + _atomic_write(snapshot_path, json.dumps(payload, indent=2) + "\n") + + log_line = {key: payload.get(key) for key in _LOG_FIELDS} + with (dataset_dir / "log.jsonl").open("a", encoding="utf-8") as fh: + fh.write(json.dumps(log_line) + "\n") + + _atomic_write(dataset_dir / "HEAD", snapshot_id + "\n") + + return SnapshotWriteResult( + dataset_id=dataset_id, + snapshot_id=snapshot_id, + snapshot_path=snapshot_path, + parent_snapshot_id=parent_snapshot_id, + is_new_dataset=is_new_dataset, + ) + + +def _check_dataset_exists(datasets_dir: Path, dataset_id: str) -> None: + """Raise unless ``dataset_id`` names a dataset already in the store. + + Args: + datasets_dir: The store's ``datasets/`` directory. Need not exist. + dataset_id: The dataset id the caller asked to continue. + + Raises: + DatasetNotFoundError: If the store holds no such dataset. + """ + if (datasets_dir / dataset_id).is_dir(): + return + known: list[str] = [] + if datasets_dir.is_dir(): + known = sorted(d.name for d in datasets_dir.iterdir() if d.is_dir()) + listing = "\n".join(f" {d}" for d in known) or " (none)" + raise DatasetNotFoundError( + f"No dataset {dataset_id!r} in the store. Omit --dataset-id to match " + f"the dataset by path, or start a new lineage.\nKnown datasets:\n{listing}" + ) + + +def _match_dataset_by_path(datasets_dir: Path, dataset_path: str) -> str | None: + """Return the id of the dataset whose latest snapshot has this path. + + Compares ``dataset_path`` against the last ``log.jsonl`` entry of every + existing dataset, the cheap lightweight record rather than the full + snapshot JSON. + + Args: + datasets_dir: The store's ``datasets/`` directory. Need not exist. + dataset_path: Relative POSIX path recorded for the snapshot being + written, as computed against the path root. + + Returns: + The ``dataset_id`` of the first dataset whose latest snapshot recorded + the same path, or ``None`` if the store holds no dataset that matches. + + Raises: + StoreCorruptionError: If a dataset's ``log.jsonl`` ends in a line that + is not valid JSON. + """ + if not datasets_dir.is_dir(): + return None + for dataset_dir in sorted(datasets_dir.iterdir()): + if not dataset_dir.is_dir(): + continue + entry = _read_head_log_entry(dataset_dir) + if entry is not None and entry.get("dataset_path") == dataset_path: + return dataset_dir.name + return None + + +def _read_head_log_entry(dataset_dir: Path) -> dict[str, Any] | None: + """Return the last (latest) parsed line of a dataset's ``log.jsonl``. + + The last line always corresponds to ``HEAD``: both are written by the same + [write_snapshot][dstrack.store.write_snapshot] call. + + Args: + dataset_dir: A ``datasets//`` directory. + + Returns: + The parsed last non-blank line, or ``None`` if the dataset has no + ``log.jsonl`` or it holds no entries yet. + + Raises: + StoreCorruptionError: If that last line is not valid JSON. + """ + log_path = dataset_dir / "log.jsonl" + if not log_path.is_file(): + return None + last_line = "" + with log_path.open(encoding="utf-8") as fh: + for line in fh: + if line.strip(): + last_line = line + if not last_line: + return None + try: + entry: dict[str, Any] = json.loads(last_line) + except json.JSONDecodeError as e: + raise StoreCorruptionError( + f"The last entry of `{log_path}` is not valid JSON. The file may " + "have been truncated by an interrupted write; restore it from git " + "or delete the trailing line." + ) from e + return entry + + +def _read_head(dataset_dir: Path) -> str | None: + """Return the snapshot_id recorded in a dataset's ``HEAD``. + + Args: + dataset_dir: A ``datasets//`` directory. + + Returns: + The latest snapshot's id, or ``None`` if the dataset has no ``HEAD`` + yet, i.e. it has never been snapshotted. + """ + head_path = dataset_dir / "HEAD" + if not head_path.is_file(): + return None + head = head_path.read_text(encoding="utf-8").strip() + return head or None + + +def _atomic_write(path: Path, text: str) -> None: + """Write ``text`` to ``path`` atomically via a temp file and ``os.replace``. + + The temp file is created in the destination directory so the final rename + stays on the same filesystem and is therefore atomic. Readers only ever see + the old content or the new one, never a partial write. + + Args: + path: File to create or replace. Its parent directory must exist. + text: Full content to write, encoded as UTF-8. + + Raises: + OSError: If the temp file cannot be written or renamed into place. The + temp file is removed and ``path`` is left untouched. + """ + fd, tmp_name = tempfile.mkstemp( + dir=path.parent, prefix=f".{path.name}.", suffix=".tmp" + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write(text) + os.replace(tmp_name, path) + except BaseException: + Path(tmp_name).unlink(missing_ok=True) + raise diff --git a/tests/test_reader_registry.py b/tests/test_reader_registry.py new file mode 100644 index 0000000..21c2f2f --- /dev/null +++ b/tests/test_reader_registry.py @@ -0,0 +1,366 @@ +"""Tests for reader registration and resolution. + +Covers `dstrack.readers._registry` and `dstrack.readers._resolve`: the three ways +a reader is reached by name (extension inference, short name, import spec), and +the protocol checks that must happen *before* a resolved class is instantiated. +""" + +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from dstrack.readers import ( + Cell, + ColumnInfo, + CsvReader, + ReaderFactory, + TabularReader, + _registry, + available_readers, + known_extensions, + load_reader_class, + register_reader, + resolve_reader, + resolve_reader_class, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +# Set by GoodReader.__init__/from_path so tests can assert whether a resolved +# class was constructed. A class that fails validation must never be. +CONSTRUCTED: list[Path] = [] + + +class GoodReader: + """A complete reader: satisfies TabularReader and ReaderFactory.""" + + EXTENSIONS = (".good",) + + @classmethod + def from_path(cls, path: Path) -> "GoodReader": + return cls(path) + + def __init__(self, path: Path) -> None: + CONSTRUCTED.append(Path(path)) + self._path = Path(path) + + def columns(self) -> list[ColumnInfo]: + return [ColumnInfo("x", "int64")] + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + return iter([[[1]]]) + + +class NoFromPathReader: + """Reads fine, but cannot be built from a path alone.""" + + def __init__(self, path: Path) -> None: + CONSTRUCTED.append(Path(path)) + + def columns(self) -> list[ColumnInfo]: + return [] + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + return iter([]) + + +class NotAReader: + """Has from_path, but none of the reading methods.""" + + @classmethod + def from_path(cls, path: Path) -> "NotAReader": + CONSTRUCTED.append(Path(path)) + return cls() + + +not_a_class = "definitely not a class" + + +class FakeEntryPoint: + """Stand-in for importlib.metadata.EntryPoint.""" + + def __init__(self, name: str, value: str, loaded: object) -> None: + self.name = name + self.value = value + self._loaded = loaded + + def load(self) -> object: + if isinstance(self._loaded, Exception): + raise self._loaded + return self._loaded + + +@pytest.fixture(autouse=True) +def isolate_registry() -> Iterator[None]: + """Restore the process-wide registry after each test.""" + by_name = dict(_registry._BY_NAME) + by_extension = dict(_registry._BY_EXTENSION) + plugins_loaded = _registry._plugins_loaded + CONSTRUCTED.clear() + yield + _registry._BY_NAME.clear() + _registry._BY_NAME.update(by_name) + _registry._BY_EXTENSION.clear() + _registry._BY_EXTENSION.update(by_extension) + _registry._plugins_loaded = plugins_loaded + + +def spec_for(cls: type) -> str: + """The 'module:ClassName' spec that resolves back to *cls*.""" + return f"{cls.__module__}:{cls.__qualname__}" + + +# --------------------------------------------------------------------------- +# Built-in registration +# --------------------------------------------------------------------------- + + +def test_csv_reader_registered_by_name_and_extension() -> None: + """The built-in CsvReader is registered out of the box, both ways.""" + assert available_readers()["csv"] is CsvReader + assert known_extensions()[".csv"] is CsvReader + + +def test_csv_reader_satisfies_reader_factory() -> None: + """ReaderFactory is checked against the class object, not an instance.""" + assert isinstance(CsvReader, ReaderFactory) + + +def test_csv_from_path_builds_a_reader(tmp_path: Path) -> None: + """from_path() constructs a working reader with default options.""" + path = tmp_path / "a.csv" + path.write_text("x\n1\n", encoding="utf-8") + + reader = CsvReader.from_path(path) + + assert isinstance(reader, TabularReader) + assert [c.name for c in reader.columns()] == ["x"] + + +# --------------------------------------------------------------------------- +# Resolution: extension, name, spec +# --------------------------------------------------------------------------- + + +def test_resolves_from_extension(tmp_path: Path) -> None: + """With no reader named, the extension picks one.""" + path = tmp_path / "a.csv" + path.write_text("x\n1\n", encoding="utf-8") + + assert isinstance(resolve_reader(path), CsvReader) + + +def test_resolves_from_short_name(tmp_path: Path) -> None: + """A registered name wins over the extension, which may be absent or wrong.""" + path = tmp_path / "export.data" + path.write_text("x\n1\n", encoding="utf-8") + + assert isinstance(resolve_reader(path, reader="csv"), CsvReader) + + +def test_resolves_from_import_spec(tmp_path: Path) -> None: + """A 'module:ClassName' spec imports a reader that is not registered at all.""" + path = tmp_path / "a.thing" + path.touch() + + reader = resolve_reader(path, reader=spec_for(GoodReader)) + + assert isinstance(reader, GoodReader) + assert [path] == CONSTRUCTED + + +def test_spec_is_told_apart_from_name_by_the_colon(tmp_path: Path) -> None: + """A value containing ':' is a spec; one without it is a registered name.""" + path = tmp_path / "a.thing" + path.touch() + + assert resolve_reader_class(path, reader=spec_for(GoodReader)) is GoodReader + + register_reader(GoodReader, name="good") + assert resolve_reader_class(path, reader="good") is GoodReader + + +def test_extension_inference_finds_a_registered_reader(tmp_path: Path) -> None: + """Registering a reader makes its extensions resolve with nothing typed. + + This is the case that --reader used to have to cover on every invocation. + """ + register_reader(GoodReader, name="good") + path = tmp_path / "data.good" + path.touch() + + assert isinstance(resolve_reader(path), GoodReader) + + +# --------------------------------------------------------------------------- +# Resolution failures +# --------------------------------------------------------------------------- + + +def test_unknown_extension_lists_the_known_ones(tmp_path: Path) -> None: + """An unclaimed extension names the ones that would have worked.""" + path = tmp_path / "data.parquet" + path.touch() + + with pytest.raises(ValueError, match=r"No reader registered for extension"): + resolve_reader(path) + + +def test_unknown_name_lists_available_readers(tmp_path: Path) -> None: + """A --reader typo names the readers that are actually registered.""" + path = tmp_path / "a.csv" + path.touch() + + with pytest.raises(ValueError, match=r"Unknown reader 'parquet'.*csv"): + resolve_reader(path, reader="parquet") + + +@pytest.mark.parametrize("spec", ["nocolon", ":Class", "module:", ""]) +def test_malformed_spec_raises_value_error(spec: str) -> None: + """A spec missing either half of 'module:ClassName' is rejected as a spec.""" + with pytest.raises(ValueError, match="Invalid reader spec"): + load_reader_class(spec) + + +def test_unimportable_module_raises_value_error() -> None: + """A spec naming a package that isn't installed fails on the import, not later.""" + with pytest.raises(ValueError, match="Could not import module"): + load_reader_class("dstrack_no_such_package:Reader") + + +def test_missing_class_raises_value_error() -> None: + """A spec whose module imports but has no such class says so.""" + with pytest.raises(ValueError, match="has no attribute"): + load_reader_class("dstrack.readers:NoSuchReader") + + +# --------------------------------------------------------------------------- +# Validation happens before instantiation +# --------------------------------------------------------------------------- + + +def test_class_missing_read_methods_is_rejected_unconstructed(tmp_path: Path) -> None: + """A class that isn't a TabularReader is rejected before it is ever called.""" + path = tmp_path / "a.thing" + path.touch() + + with pytest.raises(TypeError, match="TabularReader"): + resolve_reader(path, reader=spec_for(NotAReader)) + + assert CONSTRUCTED == [] + + +def test_class_missing_from_path_is_rejected_unconstructed(tmp_path: Path) -> None: + """Reading methods alone are not enough to be reachable by name.""" + path = tmp_path / "a.thing" + path.touch() + + with pytest.raises(TypeError, match="ReaderFactory"): + resolve_reader(path, reader=spec_for(NoFromPathReader)) + + assert CONSTRUCTED == [] + + +def test_non_class_target_is_rejected() -> None: + """A spec that names a module-level value, not a class, fails cleanly.""" + with pytest.raises(TypeError, match="not a class"): + load_reader_class(spec_for.__module__ + ":not_a_class") + + +def test_register_reader_rejects_invalid_class() -> None: + """A class that fails validation leaves no half-finished registration behind.""" + with pytest.raises(TypeError, match="ReaderFactory"): + register_reader(NoFromPathReader, name="bad") # type: ignore[arg-type] + + assert "bad" not in available_readers() + + +# --------------------------------------------------------------------------- +# Registration conflicts +# --------------------------------------------------------------------------- + + +def test_duplicate_name_raises() -> None: + """Registration never silently displaces an existing reader.""" + with pytest.raises(ValueError, match="already registered"): + register_reader(GoodReader, name="csv", extensions=()) + + +def test_duplicate_extension_raises() -> None: + """Two readers fighting over one extension is a conflict the user has to see.""" + with pytest.raises(ValueError, match=r"Extension '.csv' is already registered"): + register_reader(GoodReader, name="good", extensions=[".csv"]) + + +def test_explicit_extensions_override_the_class_attribute() -> None: + """Passing extensions replaces EXTENSIONS rather than adding to it.""" + register_reader(GoodReader, name="good", extensions=[".other"]) + + assert ".other" in known_extensions() + assert ".good" not in known_extensions() + + +# --------------------------------------------------------------------------- +# Plugin discovery +# --------------------------------------------------------------------------- + + +def test_entry_point_plugin_is_registered_lazily( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An installed plugin claims its extensions with nothing typed by the user.""" + monkeypatch.setattr( + _registry, + "entry_points", + lambda group: [FakeEntryPoint("good", spec_for(GoodReader), GoodReader)], + ) + _registry._plugins_loaded = False + + path = tmp_path / "data.good" + path.touch() + + assert isinstance(resolve_reader(path), GoodReader) + assert available_readers()["good"] is GoodReader + + +def test_broken_plugin_is_skipped_not_fatal( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """One broken plugin must not break tracking a dataset that doesn't use it.""" + monkeypatch.setattr( + _registry, + "entry_points", + lambda group: [ + FakeEntryPoint("boom", "broken:Reader", ImportError("no such module")), + FakeEntryPoint("alsobad", spec_for(NotAReader), NotAReader), + ], + ) + _registry._plugins_loaded = False + + path = tmp_path / "a.csv" + path.write_text("x\n1\n", encoding="utf-8") + + assert isinstance(resolve_reader(path), CsvReader) + assert "boom" not in available_readers() + assert "alsobad" not in available_readers() + + +def test_plugins_are_loaded_only_once(monkeypatch: pytest.MonkeyPatch) -> None: + """A failing plugin is not retried on every lookup.""" + calls = [] + + def fake_entry_points(group: str) -> list[FakeEntryPoint]: + calls.append(group) + return [] + + monkeypatch.setattr(_registry, "entry_points", fake_entry_points) + _registry._plugins_loaded = False + + available_readers() + available_readers() + known_extensions() + + assert calls == ["dstrack.readers"] diff --git a/tests/test_snapshot_metadata.py b/tests/test_snapshot_metadata.py index d6ac741..7d61e25 100644 --- a/tests/test_snapshot_metadata.py +++ b/tests/test_snapshot_metadata.py @@ -28,6 +28,7 @@ def _build( reader: TabularReader, *, path: str | Path = "/data/ds.csv", + source: str | Path | None = None, source_hash: str | None = None, ) -> SnapshotMetadata: return MetadataBuilder().build( @@ -36,6 +37,7 @@ def _build( dataset_path=path, source_type="file", created_by="tester", + source=source, source_hash=source_hash, ) @@ -181,6 +183,18 @@ def test_source_hash_empty_for_nonexistent_path() -> None: assert meta.source_hash == "" +def test_source_hashed_while_dataset_path_recorded(tmp_path: Path) -> None: + """The source is hashed; the unopenable dataset_path is recorded verbatim.""" + content = b"col\n1\n2\n" + p = tmp_path / "ds.csv" + p.write_bytes(content) + + meta = _build(_make_reader(), path="data/ds.csv", source=p) + + assert meta.dataset_path == "data/ds.csv" + assert meta.source_hash == hashlib.sha256(content).hexdigest() + + # --------------------------------------------------------------------------- # Each build produces a unique snapshot_id # --------------------------------------------------------------------------- diff --git a/tests/test_store.py b/tests/test_store.py new file mode 100644 index 0000000..a706c9f --- /dev/null +++ b/tests/test_store.py @@ -0,0 +1,99 @@ +"""Tests for write_snapshot and lineage resolution.""" + +import json +import uuid +from pathlib import Path +from typing import Any + +import pytest + +from dstrack.errors import DatasetNotFoundError +from dstrack.store import write_snapshot + + +def _snapshot(dataset_path: str = "data/ds.csv", **overrides: Any) -> dict[str, Any]: + snapshot = { + "snapshot_id": str(uuid.uuid4()), + "dataset_path": dataset_path, + "dataset_name": "ds", + "num_rows": 3, + "num_columns": 2, + } + snapshot.update(overrides) + return snapshot + + +def test_first_snapshot_mints_a_dataset(tmp_path: Path) -> None: + """A path no dataset has recorded starts a new lineage with no parent.""" + result = write_snapshot(_snapshot(), store_root=tmp_path) + + assert result.is_new_dataset + assert result.parent_snapshot_id is None + assert result.snapshot_path.is_file() + + +def test_same_path_continues_the_lineage(tmp_path: Path) -> None: + """Re-tracking the same path joins the existing dataset, parented on HEAD.""" + first = write_snapshot(_snapshot(), store_root=tmp_path) + second = write_snapshot(_snapshot(), store_root=tmp_path) + + assert not second.is_new_dataset + assert second.dataset_id == first.dataset_id + assert second.parent_snapshot_id == first.snapshot_id + + +def test_different_path_mints_a_separate_dataset(tmp_path: Path) -> None: + """An unmatched path is a new dataset, not a continuation of another.""" + first = write_snapshot(_snapshot("data/a.csv"), store_root=tmp_path) + second = write_snapshot(_snapshot("data/b.csv"), store_root=tmp_path) + + assert second.is_new_dataset + assert second.dataset_id != first.dataset_id + + +def test_explicit_dataset_id_continues_that_lineage(tmp_path: Path) -> None: + """An explicit id continues its dataset even when the path has changed.""" + first = write_snapshot(_snapshot("data/old.csv"), store_root=tmp_path) + moved = write_snapshot( + _snapshot("data/new.csv"), store_root=tmp_path, dataset_id=first.dataset_id + ) + + assert not moved.is_new_dataset + assert moved.dataset_id == first.dataset_id + assert moved.parent_snapshot_id == first.snapshot_id + + +def test_unknown_dataset_id_is_rejected(tmp_path: Path) -> None: + """An id naming no dataset fails instead of silently minting one.""" + write_snapshot(_snapshot(), store_root=tmp_path) + + with pytest.raises(DatasetNotFoundError): + write_snapshot(_snapshot(), store_root=tmp_path, dataset_id="not-a-real-id") + + +def test_rejected_dataset_id_writes_nothing(tmp_path: Path) -> None: + """A rejected id leaves no orphan dataset directory behind.""" + first = write_snapshot(_snapshot(), store_root=tmp_path) + + with pytest.raises(DatasetNotFoundError): + write_snapshot(_snapshot(), store_root=tmp_path, dataset_id="not-a-real-id") + + datasets = sorted(p.name for p in (tmp_path / "datasets").iterdir()) + assert datasets == [first.dataset_id] + + +def test_log_line_appended_per_snapshot(tmp_path: Path) -> None: + """Every snapshot appends exactly one identity line to the dataset's log.""" + first = write_snapshot(_snapshot(), store_root=tmp_path) + second = write_snapshot(_snapshot(), store_root=tmp_path) + + log_path = tmp_path / "datasets" / first.dataset_id / "log.jsonl" + entries = [json.loads(line) for line in log_path.read_text().splitlines()] + + assert [e["snapshot_id"] for e in entries] == [ + first.snapshot_id, + second.snapshot_id, + ] + assert (tmp_path / "datasets" / first.dataset_id / "HEAD").read_text().strip() == ( + second.snapshot_id + ) From 41e5dc827bc2aeaa6fd980564da521b8ab2e2164 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Tue, 14 Jul 2026 20:53:19 +0200 Subject: [PATCH 04/11] Fixes tests failind due to path interpretation in Windows vs linux based systems. --- src/dstrack/snapshot/_builder.py | 9 +++++---- src/dstrack/snapshot/_metadata.py | 14 ++++++++++---- tests/test_snapshot_metadata.py | 10 ++++++++-- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/src/dstrack/snapshot/_builder.py b/src/dstrack/snapshot/_builder.py index 7d5e562..4d9f987 100644 --- a/src/dstrack/snapshot/_builder.py +++ b/src/dstrack/snapshot/_builder.py @@ -1,5 +1,5 @@ import dataclasses -from pathlib import Path +from pathlib import Path, PurePath from typing import Any from dstrack.readers import TabularReader @@ -50,7 +50,7 @@ def build( reader: TabularReader, *, dataset_name: str, - dataset_path: str | Path, + dataset_path: str | PurePath, created_by: str, source_type: str = "file", source: str | Path | None = None, @@ -64,8 +64,9 @@ def build( Args: reader: Any [TabularReader][dstrack.readers._protocol.TabularReader]. dataset_name: Human-readable dataset name stored in the snapshot. - dataset_path: Source path or URI recorded in the snapshot. Never - opened; pass ``source`` when the data lives elsewhere. + dataset_path: Source path or URI recorded in the snapshot as a + forward-slash string. Never opened; pass ``source`` when the + data lives elsewhere. created_by: User or process identifier. source_type: Origin kind (``"file"``, ``"directory"``, etc.). source: Location the data actually lives at, used only to compute diff --git a/src/dstrack/snapshot/_metadata.py b/src/dstrack/snapshot/_metadata.py index e9c349a..d0f27d3 100644 --- a/src/dstrack/snapshot/_metadata.py +++ b/src/dstrack/snapshot/_metadata.py @@ -2,7 +2,7 @@ import uuid from dataclasses import dataclass from datetime import UTC, datetime -from pathlib import Path +from pathlib import Path, PurePath from typing import Any from dstrack.readers import ColumnInfo, TabularReader @@ -47,7 +47,7 @@ def build( reader: TabularReader, *, dataset_name: str, - dataset_path: str | Path, + dataset_path: str | PurePath, source_type: str, created_by: str, source: str | Path | None = None, @@ -64,7 +64,7 @@ def build( reader: Any TabularReader; only ``columns()`` is called. dataset_name: Human-readable dataset name stored in the snapshot. dataset_path: Source path or URI at snapshot time, recorded in the - snapshot verbatim. Never opened. + snapshot as a forward-slash string. Never opened. source_type: Origin kind (``"file"``, ``"directory"``, etc.). created_by: User or process identifier. source: Location the data actually lives at, used only to compute @@ -78,7 +78,13 @@ def build( A populated :class:`SnapshotMetadata` instance. """ cols = reader.columns() - path_str = str(dataset_path) + # Recorded with forward slashes so a store written on Windows still + # matches the same dataset when read on POSIX, and vice versa. + path_str = ( + dataset_path.as_posix() + if isinstance(dataset_path, PurePath) + else str(dataset_path) + ) if source_hash is None: p = Path(source if source is not None else dataset_path) diff --git a/tests/test_snapshot_metadata.py b/tests/test_snapshot_metadata.py index 7d61e25..7a90c2b 100644 --- a/tests/test_snapshot_metadata.py +++ b/tests/test_snapshot_metadata.py @@ -3,7 +3,7 @@ import hashlib import uuid from collections.abc import Iterator -from pathlib import Path +from pathlib import Path, PurePath, PureWindowsPath from dstrack.readers import ColumnInfo, TabularReader from dstrack.snapshot import MetadataBuilder, SnapshotMetadata @@ -27,7 +27,7 @@ def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[object]]]: def _build( reader: TabularReader, *, - path: str | Path = "/data/ds.csv", + path: str | PurePath = "/data/ds.csv", source: str | Path | None = None, source_hash: str | None = None, ) -> SnapshotMetadata: @@ -83,6 +83,12 @@ def test_dataset_path_stored_as_string() -> None: assert meta.dataset_path == "/some/path/data.csv" +def test_dataset_path_uses_forward_slashes() -> None: + """A Windows-flavoured Path is recorded with forward slashes.""" + meta = _build(_make_reader(), path=PureWindowsPath(r"some\path\data.csv")) + assert meta.dataset_path == "some/path/data.csv" + + # --------------------------------------------------------------------------- # Column descriptors # --------------------------------------------------------------------------- From e07788a971783309fd8dc53c035023bddf8439cf Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Tue, 14 Jul 2026 21:21:14 +0200 Subject: [PATCH 05/11] Fixes cross references in docstrings of package to make them compatible with zensical. --- docs/API/snapshot.md | 5 +++ src/dstrack/_benchmark/__init__.py | 10 +++--- src/dstrack/_benchmark/_environment.py | 7 +++-- src/dstrack/_benchmark/_profiling.py | 10 +++--- src/dstrack/_benchmark/_report.py | 4 +-- src/dstrack/_benchmark/_runner.py | 7 +++-- src/dstrack/_benchmark/_synthetic.py | 15 ++++++--- src/dstrack/_cli.py | 6 ++-- src/dstrack/_logging.py | 17 ++++++----- src/dstrack/paths.py | 4 +-- src/dstrack/readers/__init__.py | 6 ++-- src/dstrack/readers/_csv.py | 42 +++++++++++++++----------- src/dstrack/readers/_protocol.py | 38 ++++++++++++++--------- src/dstrack/readers/_registry.py | 25 +++++++++------ src/dstrack/readers/_resolve.py | 10 +++--- src/dstrack/snapshot/__init__.py | 10 ++++-- src/dstrack/snapshot/_metadata.py | 8 +++-- src/dstrack/snapshot/_stats.py | 11 ++++--- src/dstrack/snapshot/_version.py | 5 +-- src/dstrack/store.py | 2 +- 20 files changed, 146 insertions(+), 96 deletions(-) create mode 100644 docs/API/snapshot.md diff --git a/docs/API/snapshot.md b/docs/API/snapshot.md new file mode 100644 index 0000000..12208e6 --- /dev/null +++ b/docs/API/snapshot.md @@ -0,0 +1,5 @@ +--- +icon: lucide/camera +--- + +::: dstrack.snapshot diff --git a/src/dstrack/_benchmark/__init__.py b/src/dstrack/_benchmark/__init__.py index 5161550..b09ad3a 100644 --- a/src/dstrack/_benchmark/__init__.py +++ b/src/dstrack/_benchmark/__init__.py @@ -6,19 +6,19 @@ Layers ------ -:mod:`._synthetic` +[dstrack._benchmark._synthetic][] Describes and writes the synthetic dataset. -:mod:`._runner` +[dstrack._benchmark._runner][] Runs the snapshot pipeline against it and times each phase. -:mod:`._profiling` +[dstrack._benchmark._profiling][] Reduces the ``cProfile`` run to a call tree of dstrack's own methods. -:mod:`._environment` +[dstrack._benchmark._environment][] Describes the machine the benchmark ran on. -:mod:`._report`, :mod:`._cli` +[dstrack._benchmark._report][], [dstrack._benchmark._cli][] Render runs to a console, and wire the whole thing to ``dstrack-benchmark``. """ diff --git a/src/dstrack/_benchmark/_environment.py b/src/dstrack/_benchmark/_environment.py index 02f84eb..64bb597 100644 --- a/src/dstrack/_benchmark/_environment.py +++ b/src/dstrack/_benchmark/_environment.py @@ -1,9 +1,10 @@ """Describes the machine and runtime a benchmark executed under. Hardware facts have no portable API, so each is read by a per-OS probe. Every -probe is best-effort: it returns ``None`` (or raises, which :func:`_attempt` -swallows) when the platform will not give up the answer, and the caller falls -back to something always available. +probe is best-effort: it returns ``None`` (or raises, which +[_attempt()][dstrack._benchmark._environment._attempt] swallows) when the +platform will not give up the answer, and the caller falls back to something +always available. """ import os diff --git a/src/dstrack/_benchmark/_profiling.py b/src/dstrack/_benchmark/_profiling.py index 15b900e..03a21b1 100644 --- a/src/dstrack/_benchmark/_profiling.py +++ b/src/dstrack/_benchmark/_profiling.py @@ -1,8 +1,8 @@ """Reduces a ``cProfile`` run to a call tree of dstrack's own methods. This module is pure model: it turns raw ``pstats`` frames into -:class:`CallNode` trees and knows nothing about how they are displayed. -Rendering lives in :mod:`dstrack._benchmark._report`. +[CallNode][dstrack._benchmark._profiling.CallNode] trees and knows nothing about +how they are displayed. Rendering lives in [dstrack._benchmark._report][]. """ import cProfile @@ -91,9 +91,9 @@ class CallGraph: ``pstats`` records, for every frame, the frames that called it. This inverts those edges so a method that calls other in-scope methods (e.g. - ``StatsComputer.compute`` calling ``_build_dataset_stats``) becomes their - parent, then projects the result into trees rooted at the methods nothing - in scope called. + [StatsComputer.compute()][dstrack.snapshot._stats.StatsComputer.compute] + calling ``_build_dataset_stats``) becomes their parent, then projects the + result into trees rooted at the methods nothing in scope called. """ _scope: CallScope diff --git a/src/dstrack/_benchmark/_report.py b/src/dstrack/_benchmark/_report.py index ada0e3d..125f4a2 100644 --- a/src/dstrack/_benchmark/_report.py +++ b/src/dstrack/_benchmark/_report.py @@ -1,8 +1,8 @@ """Renders benchmark runs to a Rich console. This is the only layer that knows what a benchmark looks like on screen. It -implements :class:`~dstrack._benchmark._runner.BenchmarkObserver`, so the same -object reports progress during a run and the tables after it. +implements [BenchmarkObserver][dstrack._benchmark._runner.BenchmarkObserver], so +the same object reports progress during a run and the tables after it. """ from pathlib import Path diff --git a/src/dstrack/_benchmark/_runner.py b/src/dstrack/_benchmark/_runner.py index d10e384..ae0622d 100644 --- a/src/dstrack/_benchmark/_runner.py +++ b/src/dstrack/_benchmark/_runner.py @@ -1,9 +1,10 @@ """Runs the snapshot pipeline against a synthetic CSV and times each phase. The runner performs no output of its own: progress is announced through a -:class:`BenchmarkObserver` and everything measured is returned as a -:class:`BenchmarkRun`, so the same run can be driven from a test with no -console attached. +[BenchmarkObserver][dstrack._benchmark._runner.BenchmarkObserver] and everything +measured is returned as a +[BenchmarkRun][dstrack._benchmark._runner.BenchmarkRun], so the same run can be +driven from a test with no console attached. """ import cProfile diff --git a/src/dstrack/_benchmark/_synthetic.py b/src/dstrack/_benchmark/_synthetic.py index 8a082f6..7144ccf 100644 --- a/src/dstrack/_benchmark/_synthetic.py +++ b/src/dstrack/_benchmark/_synthetic.py @@ -1,10 +1,15 @@ """Synthetic CSV generation for the benchmark. -The dataset is described by a :class:`SyntheticCsvSpec` value object and -written by :func:`write_synthetic_csv`. Each generated dtype is one -:class:`SyntheticColumn` whose ``generate`` callable renders a single cell, so -supporting a new dtype means adding one cell factory and one entry to -``_COLUMN_KINDS`` rather than threading another count through every layer. +The dataset is described by a +[SyntheticCsvSpec][dstrack._benchmark._synthetic.SyntheticCsvSpec] value object +and written by +[write_synthetic_csv()][dstrack._benchmark._synthetic.write_synthetic_csv]. Each +generated dtype is one +[SyntheticColumn][dstrack._benchmark._synthetic.SyntheticColumn] whose +[generate][dstrack._benchmark._synthetic.SyntheticColumn.generate] callable +renders a single cell, so supporting a new dtype means adding one cell factory +and one entry to ``_COLUMN_KINDS`` rather than threading another count through +every layer. """ import csv diff --git a/src/dstrack/_cli.py b/src/dstrack/_cli.py index fe979fa..070db64 100644 --- a/src/dstrack/_cli.py +++ b/src/dstrack/_cli.py @@ -25,9 +25,9 @@ def init_local_store() -> Path: """Initializes local store structure. - Builds `STORE_TEMPLATE` (see `_store_template.py`) in a temporary - directory next to the destination, and only moves it into place once - every file and directory in the template has been created successfully. + Builds [STORE_TEMPLATE][dstrack._store_template.STORE_TEMPLATE] in a + temporary directory next to the destination, and only moves it into place + once every file and directory in the template has been created successfully. The destination is therefore never touched by a failed or partial build. Returns: diff --git a/src/dstrack/_logging.py b/src/dstrack/_logging.py index b32f615..e800b51 100644 --- a/src/dstrack/_logging.py +++ b/src/dstrack/_logging.py @@ -6,7 +6,8 @@ hierarchy however they like: ``logging.basicConfig()``, attaching handlers to the root logger, or attaching handlers directly to ``logging.getLogger("dstrack")``. -If the application does none of that, :class:`_DefaultHandler` prints +If the application does none of that, +[_DefaultHandler][dstrack._logging._DefaultHandler] prints ``WARNING``-and-above records to stderr with a sensible default format so messages are never silently dropped. It steps aside automatically the moment it detects that the application has configured logging itself. @@ -40,12 +41,14 @@ def _configure_default_logging() -> None: Called once, at import time. The handler installed here checks, on every record, whether the application has since configured logging itself and - stays silent if so, see :meth:`_DefaultHandler.emit`. - - A no-op if a :class:`_DefaultHandler` is already attached: since two of - them would each treat the other as an application-supplied handler and - both step aside, re-running this (e.g. via ``importlib.reload`` under - Jupyter's autoreload) would otherwise silence dstrack's logging entirely. + stays silent if so, see + [_DefaultHandler.emit()][dstrack._logging._DefaultHandler.emit]. + + A no-op if a [_DefaultHandler][dstrack._logging._DefaultHandler] is already + attached: since two of them would each treat the other as an + application-supplied handler and both step aside, re-running this (e.g. via + ``importlib.reload`` under Jupyter's autoreload) would otherwise silence + dstrack's logging entirely. """ if any(isinstance(h, _DefaultHandler) for h in logger.handlers): return diff --git a/src/dstrack/paths.py b/src/dstrack/paths.py index 25c2592..7997cbb 100644 --- a/src/dstrack/paths.py +++ b/src/dstrack/paths.py @@ -1,4 +1,4 @@ -"""Resolution of the local store's directory, `STORE_DIRNAME/`. +"""Resolution of the local store's directory, [STORE_DIRNAME][dstrack.paths.STORE_DIRNAME]. Mirrors how git resolves `.git/`: absent an explicit override, it is found by walking up from the current directory until a `STORE_DIRNAME/` directory turns up. @@ -55,7 +55,7 @@ def _find_store_root(start: Path) -> Path: def resolve_store_root(root: Path | str | None = None) -> Path: - """Resolve the path to the local store, `STORE_DIRNAME/`. + """Resolve the path to the local store, [STORE_DIRNAME][dstrack.paths.STORE_DIRNAME]. Checked in order, the first one available wins: 1. `root`, e.g. forwarded from a CLI `--root` option: the directory diff --git a/src/dstrack/readers/__init__.py b/src/dstrack/readers/__init__.py index 0a293c7..3bb49e4 100644 --- a/src/dstrack/readers/__init__.py +++ b/src/dstrack/readers/__init__.py @@ -2,12 +2,12 @@ Built-in readers ---------------- -[CsvReader][dstrack.readers.CsvReader] - reads ``.csv`` files; no extra dependencies. +[CsvReader][dstrack.readers._csv.CsvReader] - reads ``.csv`` files; no extra dependencies. Extending --------- -Implement [TabularReader][dstrack.readers.TabularReader] on any class to create a -custom reader. That alone is enough to build a snapshot from Python: +Implement [TabularReader][dstrack.readers._protocol.TabularReader] on any class to +create a custom reader. That alone is enough to build a snapshot from Python: Examples: ```python diff --git a/src/dstrack/readers/_csv.py b/src/dstrack/readers/_csv.py index 3dc608f..5a7d809 100644 --- a/src/dstrack/readers/_csv.py +++ b/src/dstrack/readers/_csv.py @@ -131,14 +131,16 @@ def _coerce(raw: str | None, dtype: str) -> Cell: class CsvReader: """Reads a CSV file using the standard-library ``csv`` module. - Satisfies [TabularReader][dstrack.readers._protocol.TabularReader] without inheriting from - it. Column dtypes are inferred from the first ``sample_rows`` data rows; - every subsequent call to [columns()][dstrack.readers.CsvReader.columns] returns the cached result. + Satisfies [TabularReader][dstrack.readers._protocol.TabularReader] without + inheriting from it. Column dtypes are inferred from the first + ``sample_rows`` data rows; every subsequent call to + [columns()][dstrack.readers._csv.CsvReader.columns] returns the cached + result. The file's modification time and size are recorded when schema inference - runs. [iter_batches()][dstrack.readers.CsvReader.iter_batches] checks them again before reading and raises - `RuntimeError` if the file has changed, preventing silent - schema/data mismatches. + runs. [iter_batches()][dstrack.readers._csv.CsvReader.iter_batches] checks + them again before reading and raises [RuntimeError][] if the file has + changed, preventing silent schema/data mismatches. Note: Change detection relies on ``mtime_ns`` and file size reported by the @@ -148,26 +150,27 @@ class CsvReader: also preserves the file size may go undetected. If you need guaranteed detection in such environments, ensure at least one clock tick (≥ 10 ms on most systems) elapses between calling - [columns()][dstrack.readers.CsvReader.columns] and overwriting the file. + [columns()][dstrack.readers._csv.CsvReader.columns] and overwriting the + file. Args: path: Path to the CSV file. sample_rows: Number of rows to read for dtype inference. The values are read and data types for each column are inferred from them. - encoding: File encoding passed to :func:`open`. Defaults to + encoding: File encoding passed to [open][]. Defaults to ``"utf-8"``; use ``"cp1252"`` or ``"latin-1"`` for Excel exports. rename_duplicates: When ``True``, duplicate header names are made unique by appending a counter suffix (e.g. ``col``, ``col_1``, ``col_2``). Headers that already appear exactly once in the file are treated as reserved: generated suffixes will never overwrite them (e.g. ``["a", "a", "a_1"]`` → ``["a", "a_2", "a_1"]``). - When ``False`` (default), a :exc:`ValueError` is raised instead. + When ``False`` (default), a [ValueError][] is raised instead. column_dtypes: Optional mapping of column name to dtype string that overrides the inferred dtype for those columns. Only the listed columns are affected; all others are still inferred automatically. ``"bytes"`` is not a valid override (see ADR-0002); passing it - raises `ValueError`. + raises [ValueError][]. **csv_kwargs: Forwarded verbatim to [DictReader][csv.DictReader] (e.g. ``delimiter=";"``, ``quotechar="'"``). """ @@ -187,7 +190,7 @@ def from_path(cls, path: str | Path) -> "CsvReader": path: Path to the CSV file. Returns: - A `CsvReader` bound to `path`. + A [CsvReader][dstrack.readers._csv.CsvReader] bound to ``path``. """ return cls(path) @@ -227,7 +230,8 @@ def columns(self) -> list[ColumnInfo]: """Return column descriptors, inferring dtypes on the first call. Returns: - Ordered list of [ColumnInfo][dstrack.readers.ColumnInfo] objects, one per CSV field. + Ordered list of [ColumnInfo][dstrack.readers._protocol.ColumnInfo] + objects, one per CSV field. """ if self._columns is None: self._columns = self._detect_columns() @@ -237,10 +241,12 @@ def _detect_columns(self) -> list[ColumnInfo]: """Read up to ``sample_rows`` rows and infer a ColumnInfo per field. Records ``(mtime_ns, size)`` of the open file handle in - ``_file_stat`` for later comparison by :meth:`iter_batches`. + ``_file_stat`` for later comparison by + [iter_batches()][dstrack.readers._csv.CsvReader.iter_batches]. Returns: - Ordered list of [ColumnInfo][dstrack.readers.ColumnInfo] objects, one per CSV field. + Ordered list of [ColumnInfo][dstrack.readers._protocol.ColumnInfo] + objects, one per CSV field. """ logger.debug( f"Inferring column dtypes for {self._path} from up to {self._sample_rows} sample rows", @@ -292,12 +298,14 @@ def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: batch_size: Maximum number of rows per batch. Yields: - A list of rows; each row is a list of [Cell][dstrack.readers.Cell] values - aligned with [columns()][dstrack.readers.CsvReader.columns]. + A list of rows; each row is a list of + [Cell][dstrack.readers._protocol.Cell] values aligned with + [columns()][dstrack.readers._csv.CsvReader.columns]. Raises: RuntimeError: If the file was modified since - [columns()][dstrack.readers.CsvReader.columns] was last called. + [columns()][dstrack.readers._csv.CsvReader.columns] was last + called. """ cols = self.columns() names = [c.name for c in cols] diff --git a/src/dstrack/readers/_protocol.py b/src/dstrack/readers/_protocol.py index 2c872af..ae3230e 100644 --- a/src/dstrack/readers/_protocol.py +++ b/src/dstrack/readers/_protocol.py @@ -27,9 +27,12 @@ class ColumnInfo: class TabularReader(Protocol): """Structural protocol for tabular data sources. - Any class that exposes ``columns()`` and ``iter_batches()`` satisfies this - protocol, no inheritance required. Third-party readers for Parquet, SQL, - HuggingFace datasets, etc. only need to implement these two methods. + Any class that exposes + [columns()][dstrack.readers._protocol.TabularReader.columns] and + [iter_batches()][dstrack.readers._protocol.TabularReader.iter_batches] + satisfies this protocol, no inheritance required. Third-party readers for + Parquet, SQL, HuggingFace datasets, etc. only need to implement these two + methods. Examples: >>> class MyParquetReader: @@ -49,21 +52,24 @@ def columns(self) -> list[ColumnInfo]: should return a cached result. Returns: - Ordered list of [ColumnInfo][dstrack.readers.ColumnInfo] objects, one per column. + Ordered list of [ColumnInfo][dstrack.readers._protocol.ColumnInfo] + objects, one per column. """ ... def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: """Yield non-empty batches of rows. - Each row is a list of coerced values aligned with ``columns()``. + Each row is a list of coerced values aligned with + [columns()][dstrack.readers._protocol.TabularReader.columns]. Missing values are represented as ``None``. Args: batch_size: Maximum number of rows per batch. Yields: - A list of rows, each row being a list of [Cell][dstrack.readers.Cell] values. + A list of rows, each row being a list of + [Cell][dstrack.readers._protocol.Cell] values. """ ... @@ -72,18 +78,20 @@ def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: class ReaderFactory(Protocol): """Construction contract for readers reached *by name* rather than by instance. - [TabularReader][dstrack.readers.TabularReader] describes how a reader is - *read*, and says nothing about how one is *built*: code that already holds an - instance never needs to know. But the registry and the + [TabularReader][dstrack.readers._protocol.TabularReader] describes how a + reader is *read*, and says nothing about how one is *built*: code that + already holds an instance never needs to know. But the registry and the ``"package.module:ClassName"`` spec only ever yield a class, so they need a uniform way to turn that class into an instance given a source path. - ``from_path`` is that way, and it is checked against this protocol before the - class is ever called. + [from_path()][dstrack.readers._protocol.ReaderFactory.from_path] is that way, + and it is checked against this protocol before the class is ever called. This is deliberately a second, separate protocol: a reader used only from - Python (constructed by the caller, handed straight to ``SnapshotBuilder``) - still needs nothing beyond ``TabularReader``. Only readers that are - registered, or named on the command line, must also satisfy this one. + Python (constructed by the caller, handed straight to + [SnapshotBuilder][dstrack.snapshot._builder.SnapshotBuilder]) still needs + nothing beyond [TabularReader][dstrack.readers._protocol.TabularReader]. + Only readers that are registered, or named on the command line, must also + satisfy this one. Examples: ```python @@ -107,6 +115,6 @@ def from_path(self, path: Path) -> TabularReader: Returns: An instance satisfying - [TabularReader][dstrack.readers.TabularReader]. + [TabularReader][dstrack.readers._protocol.TabularReader]. """ ... diff --git a/src/dstrack/readers/_registry.py b/src/dstrack/readers/_registry.py index beb37eb..18989b8 100644 --- a/src/dstrack/readers/_registry.py +++ b/src/dstrack/readers/_registry.py @@ -5,18 +5,21 @@ two's work: - **From Python**, a caller who already has a reader instance hands it straight - to `SnapshotBuilder`; the registry is never consulted. + to [SnapshotBuilder][dstrack.snapshot._builder.SnapshotBuilder]; the registry + is never consulted. - **From an installed plugin package**, a reader self-registers through the ``dstrack.readers`` entry-point group, so ``dstrack track data.parquet`` infers it from the extension with nothing to type. - **From an ad-hoc class** in the user's own project, which is not installed as a plugin and so is named explicitly as ``"package.module:ClassName"`` (see - `_resolve`). + [dstrack.readers._resolve][]). A registered class must satisfy both -[TabularReader][dstrack.readers.TabularReader] (how it reads) and -[ReaderFactory][dstrack.readers.ReaderFactory] (how it is built from a path); -`check_reader_class` enforces that *before* the class is ever instantiated. +[TabularReader][dstrack.readers._protocol.TabularReader] (how it reads) and +[ReaderFactory][dstrack.readers._protocol.ReaderFactory] (how it is built from a +path); +[check_reader_class()][dstrack.readers._registry.check_reader_class] enforces +that *before* the class is ever instantiated. """ import logging @@ -59,8 +62,8 @@ def check_reader_class(obj: object, *, origin: str) -> type[TabularReader]: Raises: TypeError: If ``obj`` is not a class, or does not satisfy - [TabularReader][dstrack.readers.TabularReader] and - [ReaderFactory][dstrack.readers.ReaderFactory]. + [TabularReader][dstrack.readers._protocol.TabularReader] and + [ReaderFactory][dstrack.readers._protocol.ReaderFactory]. """ if not isinstance(obj, type): raise TypeError( @@ -90,7 +93,8 @@ def build_reader(reader_cls: type[TabularReader], path: Path) -> TabularReader: """Instantiate a validated reader class for ``path``. Args: - reader_cls: A class already validated through `check_reader_class`. + reader_cls: A class already validated through + [check_reader_class()][dstrack.readers._registry.check_reader_class]. path: Path to the dataset source. Returns: @@ -114,8 +118,9 @@ def register_reader( """Register a reader under a short name and the extensions it handles. Args: - reader_cls: The reader class. Must satisfy both `TabularReader` and - `ReaderFactory`. + reader_cls: The reader class. Must satisfy both + [TabularReader][dstrack.readers._protocol.TabularReader] and + [ReaderFactory][dstrack.readers._protocol.ReaderFactory]. name: Short name, as typed for ``--reader`` (e.g. ``"csv"``). extensions: Extensions this reader claims, leading dot included. When omitted, the class's ``EXTENSIONS`` attribute is used, so a plugin diff --git a/src/dstrack/readers/_resolve.py b/src/dstrack/readers/_resolve.py index 9ca8286..866636f 100644 --- a/src/dstrack/readers/_resolve.py +++ b/src/dstrack/readers/_resolve.py @@ -1,8 +1,10 @@ -"""Selection and loading of a :class:`TabularReader` for a given source. +"""Selection and loading of a reader for a given source. -A reader is chosen one of three ways, in order of how much the user has to type: +A [TabularReader][dstrack.readers._protocol.TabularReader] is chosen one of three +ways, in order of how much the user has to type: -1. Implicitly, from the source file's extension (``data.csv`` -> `CsvReader`). +1. Implicitly, from the source file's extension (``data.csv`` -> + [CsvReader][dstrack.readers._csv.CsvReader]). Plugin readers registered through the ``dstrack.readers`` entry-point group participate here too, so an installed reader needs nothing on the command line. 2. By short name (``"csv"``), for when the extension is absent or misleading. @@ -126,7 +128,7 @@ def resolve_reader(path: str | Path, *, reader: str | None = None) -> TabularRea Returns: An instantiated reader satisfying - [TabularReader][dstrack.readers.TabularReader]. + [TabularReader][dstrack.readers._protocol.TabularReader]. Raises: ValueError: If ``reader`` names nothing known, a spec is malformed, or diff --git a/src/dstrack/snapshot/__init__.py b/src/dstrack/snapshot/__init__.py index 3253082..fffe4b6 100644 --- a/src/dstrack/snapshot/__init__.py +++ b/src/dstrack/snapshot/__init__.py @@ -2,15 +2,19 @@ Classes ------- -:class:`MetadataBuilder` +[SnapshotBuilder][dstrack.snapshot._builder.SnapshotBuilder] + Builds a complete snapshot from a single reader, combining the two below. + +[MetadataBuilder][dstrack.snapshot._metadata.MetadataBuilder] Builds identity and structural snapshot fields from a reader's schema. -:class:`StatsComputer` +[StatsComputer][dstrack.snapshot._stats.StatsComputer] Computes per-column and dataset-level statistics in a single data pass. Result types ------------ -:class:`SnapshotMetadata`, :class:`DatasetStats` +[SnapshotMetadata][dstrack.snapshot._metadata.SnapshotMetadata], +[DatasetStats][dstrack.snapshot._stats.DatasetStats] """ from ._builder import SnapshotBuilder, build_snapshot_dict diff --git a/src/dstrack/snapshot/_metadata.py b/src/dstrack/snapshot/_metadata.py index d0f27d3..cfab986 100644 --- a/src/dstrack/snapshot/_metadata.py +++ b/src/dstrack/snapshot/_metadata.py @@ -61,7 +61,9 @@ def build( written by the CLI. Args: - reader: Any TabularReader; only ``columns()`` is called. + reader: Any [TabularReader][dstrack.readers._protocol.TabularReader]; + only [columns()][dstrack.readers._protocol.TabularReader.columns] + is called. dataset_name: Human-readable dataset name stored in the snapshot. dataset_path: Source path or URI at snapshot time, recorded in the snapshot as a forward-slash string. Never opened. @@ -75,7 +77,9 @@ def build( is computed automatically. Returns: - A populated :class:`SnapshotMetadata` instance. + A populated + [SnapshotMetadata][dstrack.snapshot._metadata.SnapshotMetadata] + instance. """ cols = reader.columns() # Recorded with forward slashes so a store written on Windows still diff --git a/src/dstrack/snapshot/_stats.py b/src/dstrack/snapshot/_stats.py index 6af0a4e..b87b77f 100644 --- a/src/dstrack/snapshot/_stats.py +++ b/src/dstrack/snapshot/_stats.py @@ -108,8 +108,9 @@ class DatasetStats: class _ColumnAcc(Protocol): """Common interface every per-column accumulator implements. - ``compute`` drives all accumulators through this interface so it does not - need to branch on column type while scanning rows or building stats. + [compute()][dstrack.snapshot._stats.StatsComputer.compute] drives all + accumulators through this interface so it does not need to branch on column + type while scanning rows or building stats. """ null_count: int @@ -232,10 +233,12 @@ def compute(self, reader: TabularReader) -> DatasetStats: """Run a full data pass and return aggregated statistics. Args: - reader: Any TabularReader whose batches will be consumed once. + reader: Any [TabularReader][dstrack.readers._protocol.TabularReader] + whose batches will be consumed once. Returns: - A populated :class:`DatasetStats` instance. + A populated [DatasetStats][dstrack.snapshot._stats.DatasetStats] + instance. """ cols = reader.columns() col_names = [c.name for c in cols] diff --git a/src/dstrack/snapshot/_version.py b/src/dstrack/snapshot/_version.py index 2a7a255..31f65b1 100644 --- a/src/dstrack/snapshot/_version.py +++ b/src/dstrack/snapshot/_version.py @@ -1,7 +1,8 @@ """Single source of truth for the snapshot format version. -Bump ``FORMAT_VERSION`` and ``SCHEMA_PATH`` together whenever -``SnapshotMetadata``'s shape changes. +Bump [FORMAT_VERSION][dstrack.snapshot._version.FORMAT_VERSION] and +[SCHEMA_PATH][dstrack.snapshot._version.SCHEMA_PATH] together whenever +[SnapshotMetadata][dstrack.snapshot._metadata.SnapshotMetadata]'s shape changes. """ from pathlib import Path diff --git a/src/dstrack/store.py b/src/dstrack/store.py index a36bb95..f028003 100644 --- a/src/dstrack/store.py +++ b/src/dstrack/store.py @@ -70,7 +70,7 @@ def write_snapshot( Args: snapshot: A snapshot dict as produced by - [SnapshotBuilder][dstrack.snapshot.SnapshotBuilder]. Must contain + [SnapshotBuilder][dstrack.snapshot._builder.SnapshotBuilder]. Must contain ``snapshot_id`` and ``dataset_path``; ``parent_snapshot_id`` is filled in here and overwritten if already present. store_root: Path to the ``.dstrack/`` directory, e.g. from From 21df1c0b8eeb9303d20b0114760adbdfa631b908 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 14:24:50 +0200 Subject: [PATCH 06/11] Sets target code coverage for project --- codecov.yaml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/codecov.yaml b/codecov.yaml index eb14bfb..5d5e3ea 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -2,9 +2,8 @@ coverage: status: project: default: - target: auto + target: 95% threshold: 0.5% - base: auto comment: layout: "diff, flags, files" behavior: default From 4613394bce9eadba7773a316f9eedf9b013eec67 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 14:53:46 +0200 Subject: [PATCH 07/11] Adds benchmark package unit tests. --- src/dstrack/_benchmark/_cli.py | 11 + tests/test_benchmark_cli.py | 386 +++++++++++++++++++++++++++++++++ 2 files changed, 397 insertions(+) create mode 100644 tests/test_benchmark_cli.py diff --git a/src/dstrack/_benchmark/_cli.py b/src/dstrack/_benchmark/_cli.py index 4c00f57..c943739 100644 --- a/src/dstrack/_benchmark/_cli.py +++ b/src/dstrack/_benchmark/_cli.py @@ -22,6 +22,17 @@ _DEFAULTS = SyntheticCsvSpec() +@app.callback() +def _main() -> None: + """Keep ``run`` a named subcommand so a bare invocation shows help. + + Typer promotes a lone command to the top level, which quietly discards + ``no_args_is_help``; an explicit callback keeps the command group intact so + running ``dstrack-benchmark`` with no arguments prints help instead of + silently kicking off a full-size benchmark. + """ + + @contextmanager def _workspace(reporter: ConsoleReporter, *, keep_csv: bool) -> Iterator[Path]: """Yield a path for the synthetic CSV, discarding it unless ``keep_csv``.""" diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py new file mode 100644 index 0000000..1ce5596 --- /dev/null +++ b/tests/test_benchmark_cli.py @@ -0,0 +1,386 @@ +"""Tests for the benchmark CLI in src/dstrack/_benchmark/_cli.py.""" + +import shutil +from dataclasses import dataclass +from pathlib import Path + +import pytest +from typer.testing import CliRunner + +from dstrack._benchmark import _cli +from dstrack._benchmark._cli import _workspace, app +from dstrack._benchmark._environment import collect_environment_info +from dstrack._benchmark._runner import BenchmarkResult, BenchmarkRun +from dstrack._benchmark._synthetic import SyntheticCsvSpec + +runner = CliRunner() + +# `run` is a named subcommand: an explicit app callback keeps the command group +# intact, so a bare `dstrack-benchmark` shows help instead of benchmarking. + + +@dataclass +class _RecordingReporter: + """Stands in for ConsoleReporter, recording what the CLI asked it to print.""" + + kept_paths: list[Path] + reported: list[tuple[BenchmarkRun, int]] + + def __init__(self) -> None: + self.kept_paths = [] + self.reported = [] + + def report(self, run: BenchmarkRun, *, profile_limit: int) -> None: + self.reported.append((run, profile_limit)) + + def csv_kept(self, path: Path) -> None: + self.kept_paths.append(path) + + +def _make_run() -> BenchmarkRun: + """A canned run, so stubbed runners need not do any real work.""" + result = BenchmarkResult( + num_rows=10, + num_rows_read=10, + num_columns=3, + csv_size_bytes=128, + generation_seconds=0.1, + metadata_seconds=0.2, + stats_seconds=0.3, + schema_hash="deadbeef", + ) + return BenchmarkRun( + result=result, environment=collect_environment_info(), call_graph=None + ) + + +class _StubRunner: + """Stands in for BenchmarkRunner, recording how the CLI constructed it.""" + + def __init__( + self, spec: SyntheticCsvSpec, *, observer: object = None, profile: bool = True + ) -> None: + self.instances: list[_StubRunner] = [] + self.spec = spec + self.observer = observer + self.profile = profile + self.csv_paths: list[Path] = [] + type(self).instances.append(self) + + def run(self, csv_path: Path) -> BenchmarkRun: + csv_path.write_text("id\n0\n", encoding="utf-8") + self.csv_paths.append(csv_path) + return _make_run() + + +@pytest.fixture +def stubbed(monkeypatch: pytest.MonkeyPatch) -> _RecordingReporter: + """Swap the runner and reporter out, so the CLI's wiring is tested alone.""" + _StubRunner.instances = [] + reporter = _RecordingReporter() + monkeypatch.setattr(_cli, "BenchmarkRunner", _StubRunner) + monkeypatch.setattr(_cli, "ConsoleReporter", lambda: reporter) + return reporter + + +# --------------------------------------------------------------------------- +# _workspace +# --------------------------------------------------------------------------- + + +def test_workspace_yields_csv_path_in_a_fresh_directory() -> None: + """Yields `synthetic_dataset.csv` inside a directory that exists but is empty.""" + reporter = _RecordingReporter() + + with _workspace(reporter, keep_csv=False) as csv_path: + assert csv_path.name == "synthetic_dataset.csv" + assert csv_path.parent.is_dir() + assert not csv_path.exists() + assert list(csv_path.parent.iterdir()) == [] + + +def test_workspace_discards_the_directory_by_default() -> None: + """Removes the temp directory, and the CSV in it, once the block exits.""" + reporter = _RecordingReporter() + + with _workspace(reporter, keep_csv=False) as csv_path: + csv_path.write_text("id\n0\n", encoding="utf-8") + tmp_dir = csv_path.parent + + assert not csv_path.exists() + assert not tmp_dir.exists() + assert reporter.kept_paths == [] + + +def test_workspace_keeps_the_csv_when_asked() -> None: + """With keep_csv, leaves the CSV on disk and reports where it went.""" + reporter = _RecordingReporter() + + with _workspace(reporter, keep_csv=True) as csv_path: + csv_path.write_text("id\n0\n", encoding="utf-8") + + assert csv_path.is_file() + assert csv_path.read_text(encoding="utf-8") == "id\n0\n" + assert reporter.kept_paths == [csv_path] + + # keep_csv deliberately leaves the directory behind, so clean it up here. + shutil.rmtree(csv_path.parent, ignore_errors=True) + + +def test_workspace_cleans_up_when_the_block_raises() -> None: + """Propagates the error, but still discards the temp directory.""" + reporter = _RecordingReporter() + leaked: list[Path] = [] + + with ( + pytest.raises(RuntimeError, match="boom"), + _workspace(reporter, keep_csv=False) as csv_path, + ): + csv_path.write_text("id\n0\n", encoding="utf-8") + leaked.append(csv_path.parent) + raise RuntimeError("boom") + + assert not leaked[0].exists() + + +def test_workspace_uses_a_distinct_directory_per_run() -> None: + """Concurrent or repeated runs never collide on the same temp directory.""" + reporter = _RecordingReporter() + + with ( + _workspace(reporter, keep_csv=False) as first, + _workspace(reporter, keep_csv=False) as second, + ): + assert first.parent != second.parent + + +# --------------------------------------------------------------------------- +# `run` command: option wiring +# --------------------------------------------------------------------------- + + +def test_run_defaults_match_the_synthetic_spec_defaults( + stubbed: _RecordingReporter, +) -> None: + """With no options, the spec the runner gets is an untouched SyntheticCsvSpec.""" + result = runner.invoke(app, ["run"]) + + assert result.exit_code == 0 + assert _StubRunner.instances[0].spec == SyntheticCsvSpec() + + +def test_run_forwards_every_option_to_the_spec(stubbed: _RecordingReporter) -> None: + """Each dataset-shaping option lands on the matching spec field.""" + result = runner.invoke( + app, + [ + "run", + "--rows", "1234", + "--numeric-cols", "7", + "--string-cols", "6", + "--datetime-cols", "5", + "--bool-cols", "4", + "--null-rate", "0.25", + "--seed", "99", + ], + ) # fmt: skip + + assert result.exit_code == 0 + assert _StubRunner.instances[0].spec == SyntheticCsvSpec( + num_rows=1234, + num_numeric_cols=7, + num_string_cols=6, + num_datetime_cols=5, + num_bool_cols=4, + null_rate=0.25, + seed=99, + ) + + +def test_run_uses_the_reporter_as_the_runners_observer( + stubbed: _RecordingReporter, +) -> None: + """The reporter that prints the tables also receives the progress callbacks.""" + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code == 0 + assert _StubRunner.instances[0].observer is stubbed + + +def test_run_profiles_by_default(stubbed: _RecordingReporter) -> None: + """Profiling is on unless it is explicitly turned off.""" + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code == 0 + assert _StubRunner.instances[0].profile is True + + +def test_run_no_profile_disables_profiling(stubbed: _RecordingReporter) -> None: + """--no-profile reaches the runner, so no profiler is attached.""" + result = runner.invoke(app, ["run", "--rows", "10", "--no-profile"]) + + assert result.exit_code == 0 + assert _StubRunner.instances[0].profile is False + + +def test_run_forwards_the_profile_limit_to_the_report( + stubbed: _RecordingReporter, +) -> None: + """--profile-limit is a reporting concern, and reaches report().""" + result = runner.invoke(app, ["run", "--rows", "10", "--profile-limit", "3"]) + + assert result.exit_code == 0 + assert [limit for _, limit in stubbed.reported] == [3] + + +def test_run_defaults_the_profile_limit_to_20(stubbed: _RecordingReporter) -> None: + """The default call-tree width is 20 children per node.""" + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code == 0 + assert [limit for _, limit in stubbed.reported] == [20] + + +def test_run_reports_the_run_the_runner_produced(stubbed: _RecordingReporter) -> None: + """Whatever the runner returns is handed straight to the reporter.""" + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code == 0 + assert len(stubbed.reported) == 1 + reported_run, _ = stubbed.reported[0] + assert reported_run.result.schema_hash == "deadbeef" + + +def test_run_benchmarks_the_csv_inside_the_workspace( + stubbed: _RecordingReporter, +) -> None: + """The path handed to the runner is the workspace CSV, and it is cleaned up.""" + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code == 0 + csv_path = _StubRunner.instances[0].csv_paths[0] + assert csv_path.name == "synthetic_dataset.csv" + assert not csv_path.parent.exists() + assert stubbed.kept_paths == [] + + +def test_run_keep_csv_leaves_the_generated_csv_behind( + stubbed: _RecordingReporter, +) -> None: + """With --keep-csv the file survives the run and is reported to the user.""" + result = runner.invoke(app, ["run", "--rows", "10", "--keep-csv"]) + + assert result.exit_code == 0 + csv_path = _StubRunner.instances[0].csv_paths[0] + assert csv_path.is_file() + assert stubbed.kept_paths == [csv_path] + + shutil.rmtree(csv_path.parent, ignore_errors=True) + + +def test_run_rejects_a_non_numeric_option(stubbed: _RecordingReporter) -> None: + """Bad option types fail as a usage error, without constructing a runner.""" + result = runner.invoke(app, ["run", "--rows", "many"]) + + assert result.exit_code != 0 + assert _StubRunner.instances == [] + + +# --------------------------------------------------------------------------- +# `run` command: end to end +# --------------------------------------------------------------------------- + + +def test_run_end_to_end_prints_all_three_tables() -> None: + """A real (tiny) run generates a CSV, snapshots it, and prints every table.""" + result = runner.invoke(app, ["run", "--rows", "50", "--profile-limit", "5"]) + + assert result.exit_code == 0 + output = result.output + assert "Generating synthetic CSV" in output + assert "Building snapshot metadata" in output + assert "Computing snapshot statistics" in output + assert "Environment" in output + assert "Snapshot benchmark" in output + assert "call tree" in output + + +def test_run_end_to_end_without_profiling_omits_the_call_tree() -> None: + """--no-profile still reports timings, but there is no call tree to show.""" + result = runner.invoke(app, ["run", "--rows", "50", "--no-profile"]) + + assert result.exit_code == 0 + assert "Snapshot benchmark" in result.output + assert "call tree" not in result.output + + +def test_run_end_to_end_keep_csv_writes_a_readable_dataset( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The kept CSV is the dataset that was benchmarked: header plus the rows asked for.""" + monkeypatch.setattr(_cli.tempfile, "mkdtemp", lambda prefix: str(tmp_path)) + + result = runner.invoke( + app, ["run", "--rows", "50", "--string-cols", "1", "--no-profile", "--keep-csv"] + ) + + assert result.exit_code == 0 + csv_path = tmp_path / "synthetic_dataset.csv" + assert csv_path.is_file() + assert "Synthetic CSV kept at" in result.output + + lines = csv_path.read_text(encoding="utf-8").splitlines() + assert lines[0].startswith("id,") + assert "string_0" in lines[0] + assert len(lines) == 51 # header + 50 rows + + +def test_run_end_to_end_is_reproducible_for_a_fixed_seed( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The same seed benchmarks the same bytes, so runs are comparable.""" + contents: list[str] = [] + for i in range(2): + out_dir = tmp_path / str(i) + out_dir.mkdir() + monkeypatch.setattr(_cli.tempfile, "mkdtemp", lambda prefix, d=out_dir: str(d)) + + result = runner.invoke( + app, ["run", "--rows", "50", "--seed", "7", "--no-profile", "--keep-csv"] + ) + + assert result.exit_code == 0 + contents.append((out_dir / "synthetic_dataset.csv").read_text(encoding="utf-8")) + + assert contents[0] == contents[1] + + +def test_run_surfaces_a_runner_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """A failure mid-run exits non-zero rather than reporting a bogus benchmark.""" + + class _ExplodingRunner: + def __init__(self, spec: SyntheticCsvSpec, **kwargs: object) -> None: + self.leaked: list[Path] = [] + + def run(self, csv_path: Path) -> BenchmarkRun: + leaked.append(csv_path.parent) + raise RuntimeError("reader exploded") + + leaked: list[Path] = [] + monkeypatch.setattr(_cli, "BenchmarkRunner", _ExplodingRunner) + + result = runner.invoke(app, ["run", "--rows", "10"]) + + assert result.exit_code != 0 + assert isinstance(result.exception, RuntimeError) + assert not leaked[0].exists() # the workspace is still cleaned up + + +def test_no_args_shows_help_without_running_a_benchmark() -> None: + """A bare invocation prints help and lists `run`, never starting a benchmark.""" + result = runner.invoke(app, []) + + assert result.exit_code != 0 + assert "Usage" in result.output + assert "run" in result.output + assert "Snapshot benchmark" not in result.output From c8e748dde866c3bde41e8a1057fe874533742b3f Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 15:22:35 +0200 Subject: [PATCH 08/11] Adds tests for _track Typer script. --- tests/test_track.py | 266 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 266 insertions(+) create mode 100644 tests/test_track.py diff --git a/tests/test_track.py b/tests/test_track.py new file mode 100644 index 0000000..a849b9c --- /dev/null +++ b/tests/test_track.py @@ -0,0 +1,266 @@ +"""Tests for the `dstrack track` command in src/dstrack/_track.py.""" + +import json +from pathlib import Path +from typing import Any + +import pytest +from typer.testing import CliRunner + +from dstrack import _track +from dstrack._cli import app +from dstrack._track import _default_created_by + +runner = CliRunner() + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture +def store(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Create an empty local store under a fresh cwd and return its root. + + The `track` command resolves the store by walking up from the current + directory, so the tests run from inside `tmp_path`. `DSTRACK_ROOT_PATH` + is cleared so an ambient value cannot redirect the store elsewhere. + """ + monkeypatch.delenv("DSTRACK_ROOT_PATH", raising=False) + monkeypatch.chdir(tmp_path) + store_root = tmp_path / ".dstrack" + (store_root / "datasets").mkdir(parents=True) + return store_root + + +def _write_csv(path: Path, rows: int = 3) -> Path: + """Write a small, valid CSV file and return its path.""" + lines = ["a,b"] + [f"{i},{i * 2}" for i in range(rows)] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +def _load_only_snapshot(store_root: Path) -> dict[str, Any]: + """Load the single snapshot JSON written under the store.""" + snapshots = list(store_root.glob("datasets/*/snapshots/*.json")) + assert len(snapshots) == 1, f"expected exactly one snapshot, got {snapshots}" + return json.loads(snapshots[0].read_text(encoding="utf-8")) + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +def test_track_creates_new_dataset_snapshot(tmp_path: Path, store: Path) -> None: + """Tracking a fresh file writes a snapshot and reports a new dataset.""" + csv = _write_csv(tmp_path / "data.csv") + + result = runner.invoke(app, ["track", str(csv)]) + + assert result.exit_code == 0, result.output + assert "new dataset" in result.output + assert "written" in result.output + snapshot = _load_only_snapshot(store) + assert snapshot["num_rows"] == 3 + assert snapshot["parent_snapshot_id"] is None + + +def test_track_same_file_twice_continues_lineage(tmp_path: Path, store: Path) -> None: + """Re-tracking the same path extends the lineage rather than starting one.""" + csv = _write_csv(tmp_path / "data.csv") + + first = runner.invoke(app, ["track", str(csv)]) + second = runner.invoke(app, ["track", str(csv)]) + + assert first.exit_code == 0, first.output + assert second.exit_code == 0, second.output + assert "new dataset" in first.output + assert "continued lineage" in second.output + # Both snapshots live under the same dataset directory. + dataset_dirs = list((store / "datasets").iterdir()) + assert len(dataset_dirs) == 1 + # Snapshot dir should contain 2 snapshots + HEAD file + snapshot_dir = dataset_dirs[0] + assert len(list(snapshot_dir.iterdir())) == 3 + + +def test_track_default_dataset_name_is_file_stem(tmp_path: Path, store: Path) -> None: + """Without --name the dataset name defaults to the file's stem.""" + csv = _write_csv(tmp_path / "customers.csv") + + result = runner.invoke(app, ["track", str(csv)]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["dataset_name"] == "customers" + + +def test_track_name_option_overrides_stem(tmp_path: Path, store: Path) -> None: + """--name overrides the default dataset name derived from the file stem.""" + csv = _write_csv(tmp_path / "customers.csv") + + result = runner.invoke(app, ["track", str(csv), "--name", "clients"]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["dataset_name"] == "clients" + + +def test_track_created_by_option_is_recorded(tmp_path: Path, store: Path) -> None: + """--created-by overrides the recorded author.""" + csv = _write_csv(tmp_path / "data.csv") + + result = runner.invoke(app, ["track", str(csv), "--created-by", "alice"]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["created_by"] == "alice" + + +def test_track_records_path_relative_to_store_root(tmp_path: Path, store: Path) -> None: + """The dataset_path is stored relative to the store's parent by default.""" + sub = tmp_path / "sub" + sub.mkdir() + csv = _write_csv(sub / "data.csv") + + result = runner.invoke(app, ["track", str(csv)]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["dataset_path"] == "sub/data.csv" + + +def test_track_root_option_changes_recorded_path(tmp_path: Path, store: Path) -> None: + """--root controls what the recorded dataset_path is made relative to.""" + sub = tmp_path / "sub" + sub.mkdir() + csv = _write_csv(sub / "data.csv") + + result = runner.invoke(app, ["track", str(csv), "--root", str(sub)]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["dataset_path"] == "data.csv" + + +def test_track_explicit_reader_reads_mislabeled_file( + tmp_path: Path, store: Path +) -> None: + """--reader wins over extension inference, so a .data file reads as CSV.""" + csv = _write_csv(tmp_path / "data.unknownext") + + result = runner.invoke(app, ["track", str(csv), "--reader", "csv"]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["num_rows"] == 3 + + +# --------------------------------------------------------------------------- +# `--dataset-id` +# --------------------------------------------------------------------------- + + +def test_track_dataset_id_continues_that_lineage(tmp_path: Path, store: Path) -> None: + """--dataset-id attaches the snapshot to an existing lineage after a move.""" + first_csv = _write_csv(tmp_path / "old.csv") + first = runner.invoke(app, ["track", str(first_csv)]) + assert first.exit_code == 0, first.output + dataset_id = next((store / "datasets").iterdir()).name + + moved_csv = _write_csv(tmp_path / "new.csv") + second = runner.invoke(app, ["track", str(moved_csv), "--dataset-id", dataset_id]) + + assert second.exit_code == 0, second.output + assert "continued lineage" in second.output + assert len(list((store / "datasets").iterdir())) == 1 + + +def test_track_unknown_dataset_id_errors(tmp_path: Path, store: Path) -> None: + """An unknown --dataset-id fails cleanly instead of minting a lineage.""" + csv = _write_csv(tmp_path / "data.csv") + + result = runner.invoke(app, ["track", str(csv), "--dataset-id", "does-not-exist"]) + + assert result.exit_code == 1 + assert "does-not-exist" in result.output + + +# --------------------------------------------------------------------------- +# Failure paths +# --------------------------------------------------------------------------- + + +def test_track_without_store_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """With no `.dstrack/` anywhere above, the command exits non-zero.""" + monkeypatch.delenv("DSTRACK_ROOT_PATH", raising=False) + monkeypatch.chdir(tmp_path) + csv = _write_csv(tmp_path / "data.csv") + + result = runner.invoke(app, ["track", str(csv)]) + + assert result.exit_code == 1 + assert ".dstrack" in result.output + + +def test_track_unknown_extension_hints_at_reader(tmp_path: Path, store: Path) -> None: + """An unresolvable extension fails and suggests naming a reader explicitly.""" + data = _write_csv(tmp_path / "data.unknownext") + + result = runner.invoke(app, ["track", str(data)]) + + assert result.exit_code == 1 + assert "--reader" in result.output + + +def test_track_unknown_named_reader_errors_without_hint( + tmp_path: Path, store: Path +) -> None: + """A bad explicit --reader fails without appending the --reader hint.""" + csv = _write_csv(tmp_path / "data.csv") + + result = runner.invoke(app, ["track", str(csv), "--reader", "nope"]) + + assert result.exit_code == 1 + assert "Unknown reader" in result.output + assert "Use --reader to name one explicitly." not in result.output + + +def test_track_missing_path_is_rejected_by_typer(tmp_path: Path, store: Path) -> None: + """A path that does not exist is rejected by argument validation.""" + result = runner.invoke(app, ["track", str(tmp_path / "missing.csv")]) + + assert result.exit_code != 0 + assert result.exit_code != 1 # a usage error, not our handled Exit(1) + + +def test_track_directory_path_is_rejected_by_typer(tmp_path: Path, store: Path) -> None: + """A directory is rejected: the argument is declared dir_okay=False.""" + result = runner.invoke(app, ["track", str(tmp_path)]) + + assert result.exit_code != 0 + + +# --------------------------------------------------------------------------- +# _default_created_by +# --------------------------------------------------------------------------- + + +def test_default_created_by_returns_username( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Returns the current username when getpass can determine one.""" + monkeypatch.setattr(_track.getpass, "getuser", lambda: "bob") + + assert _default_created_by() == "bob" + + +def test_default_created_by_falls_back_to_unknown( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Falls back to 'unknown' when the username cannot be determined.""" + + def _boom() -> str: + raise OSError("no username") + + monkeypatch.setattr(_track.getpass, "getuser", _boom) + + assert _default_created_by() == "unknown" From 72d063742d95ee13057e56246bcd4b07e9d432a2 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 18:06:10 +0200 Subject: [PATCH 09/11] Adds solutions for review comments. --- .../0001-dataset-snapshot-content.md | 1 + src/dstrack/_benchmark/_cli.py | 16 +- src/dstrack/_benchmark/_profiling.py | 17 +- src/dstrack/_benchmark/_runner.py | 5 + src/dstrack/_benchmark/_synthetic.py | 21 ++ src/dstrack/_track.py | 35 ++- src/dstrack/errors.py | 4 + src/dstrack/readers/_registry.py | 35 ++- src/dstrack/snapshot/_builder.py | 4 +- src/dstrack/snapshot/_stats.py | 53 ++++- src/dstrack/store.py | 221 +++++++++++++----- tests/test_benchmark_cli.py | 4 +- tests/test_reader_registry.py | 76 ++++++ tests/test_snapshot_stats.py | 110 +++++++++ tests/test_store.py | 45 ++++ tests/test_track.py | 24 ++ 16 files changed, 596 insertions(+), 75 deletions(-) diff --git a/docs/decisions/0001-dataset-snapshot-content.md b/docs/decisions/0001-dataset-snapshot-content.md index 03d602a..3d1507f 100644 --- a/docs/decisions/0001-dataset-snapshot-content.md +++ b/docs/decisions/0001-dataset-snapshot-content.md @@ -105,3 +105,4 @@ These compact structures enable cheap diff and similarity queries across snapsho - Content fingerprints (MinHash, HyperLogLog) enable cheap cross-snapshot diff and similarity queries without loading raw data, at the cost of being approximate. - Scope is intentionally limited to tabular data; non-tabular modalities (images, audio, video, graphs) are not supported and would require a new ADR to extend the schema. - `format_version` is required on every snapshot to allow the schema to evolve without breaking existing readers. +- Statistics are computed in a single **in-memory** pass, so worst-case memory grows with the row count (per-column values, the distinct-row set, and per-string value counts are retained until the pass completes). A configurable `max_rows` limit bounds this: datasets larger than the limit are rejected with a clear error rather than silently exhausting memory. The limit is exposed on `StatsComputer` and via the `dstrack track --max-rows` option, and may be disabled from Python when enough memory is available. diff --git a/src/dstrack/_benchmark/_cli.py b/src/dstrack/_benchmark/_cli.py index c943739..7795b55 100644 --- a/src/dstrack/_benchmark/_cli.py +++ b/src/dstrack/_benchmark/_cli.py @@ -50,22 +50,23 @@ def _workspace(reporter: ConsoleReporter, *, keep_csv: bool) -> Iterator[Path]: @app.command(help="Generate a synthetic CSV and benchmark snapshot creation on it.") def run( rows: Annotated[ - int, typer.Option(help="Number of data rows to generate.") + int, typer.Option(min=0, help="Number of data rows to generate.") ] = _DEFAULTS.num_rows, numeric_cols: Annotated[ - int, typer.Option(help="Number of numeric (float) columns.") + int, typer.Option(min=0, help="Number of numeric (float) columns.") ] = _DEFAULTS.num_numeric_cols, string_cols: Annotated[ - int, typer.Option(help="Number of string columns.") + int, typer.Option(min=0, help="Number of string columns.") ] = _DEFAULTS.num_string_cols, datetime_cols: Annotated[ - int, typer.Option(help="Number of datetime columns.") + int, typer.Option(min=0, help="Number of datetime columns.") ] = _DEFAULTS.num_datetime_cols, bool_cols: Annotated[ - int, typer.Option(help="Number of boolean columns.") + int, typer.Option(min=0, help="Number of boolean columns.") ] = _DEFAULTS.num_bool_cols, null_rate: Annotated[ - float, typer.Option(help="Fraction of non-id cells left null.") + float, + typer.Option(min=0.0, max=1.0, help="Fraction of non-id cells left null."), ] = _DEFAULTS.null_rate, seed: Annotated[ int, typer.Option(help="Random seed for reproducible synthetic data.") @@ -86,8 +87,9 @@ def run( profile_limit: Annotated[ int, typer.Option( + min=0, help="Max number of child methods shown under each node of the " - "profile call tree, ranked by cumulative time." + "profile call tree, ranked by cumulative time.", ), ] = 20, ) -> None: diff --git a/src/dstrack/_benchmark/_profiling.py b/src/dstrack/_benchmark/_profiling.py index 03a21b1..7ca9fcc 100644 --- a/src/dstrack/_benchmark/_profiling.py +++ b/src/dstrack/_benchmark/_profiling.py @@ -108,6 +108,7 @@ def from_profile(cls, profiler: cProfile.Profile, scope: CallScope) -> "CallGrap graph = cls(_scope=scope) frames = {key for key in raw if scope.contains(key[0])} callees: set[FrameKey] = set() + entrypoints: set[FrameKey] = set() for key in frames: primitive_calls, num_calls, total_time, cumulative_time, callers = raw[key] @@ -121,8 +122,14 @@ def from_profile(cls, profiler: cProfile.Profile, scope: CallScope) -> "CallGrap if caller in frames: graph._children.setdefault(caller, []).append(key) callees.add(key) - - graph._roots = graph._by_cumulative(frames - callees) + else: + entrypoints.add(key) + + # A frame is a root if nothing in scope called it, or if it was also + # entered from outside the scope. The latter keeps external entry + # points into a recursive component (f -> g -> f) as roots, which + # ``frames - callees`` alone would drop, leaving the tree empty. + graph._roots = graph._by_cumulative((frames - callees) | entrypoints) return graph @property @@ -136,7 +143,13 @@ def trees(self, *, max_children: int) -> list[CallNode]: Args: max_children: Maximum callees expanded under each node, ranked by cumulative time. The rest are counted in ``hidden_children``. + + Raises: + ValueError: If ``max_children`` is negative, which would slice the + callee list from the end rather than cap it as documented. """ + if max_children < 0: + raise ValueError(f"max_children must be >= 0, got {max_children}") return [ self._expand(root, ancestors=frozenset({root}), max_children=max_children) for root in self._roots diff --git a/src/dstrack/_benchmark/_runner.py b/src/dstrack/_benchmark/_runner.py index ae0622d..b426be2 100644 --- a/src/dstrack/_benchmark/_runner.py +++ b/src/dstrack/_benchmark/_runner.py @@ -132,6 +132,11 @@ def __init__( def run(self, csv_path: Path) -> BenchmarkRun: """Generate the dataset at ``csv_path`` and benchmark a snapshot of it.""" + # clean profiler if already created to avoid misleading profiles on re-runs. + if self._profiler is not None: + self._profiler.clear() + + # generate data for benchmark self._observer.generating_csv(csv_path, self._spec.num_rows) with self._timed() as generation: write_synthetic_csv(csv_path, self._spec) diff --git a/src/dstrack/_benchmark/_synthetic.py b/src/dstrack/_benchmark/_synthetic.py index 7144ccf..87169d8 100644 --- a/src/dstrack/_benchmark/_synthetic.py +++ b/src/dstrack/_benchmark/_synthetic.py @@ -13,6 +13,7 @@ """ import csv +import math import random from collections.abc import Callable from dataclasses import dataclass @@ -96,6 +97,26 @@ class SyntheticCsvSpec: null_rate: float = 0.02 seed: int = 42 + def __post_init__(self) -> None: + """Reject counts and a null rate that would silently corrupt the dataset. + + Negative counts and an out-of-range ``null_rate`` have no meaningful + dataset to generate, so they are refused here rather than producing a + subtly wrong CSV further downstream. + """ + for name in ( + "num_rows", + "num_numeric_cols", + "num_string_cols", + "num_datetime_cols", + "num_bool_cols", + ): + value: int = getattr(self, name) + if value < 0: + raise ValueError(f"{name} must be >= 0, got {value}") + if not (math.isfinite(self.null_rate) and 0.0 <= self.null_rate <= 1.0): + raise ValueError(f"null_rate must be in [0, 1], got {self.null_rate!r}") + def columns(self) -> list[SyntheticColumn]: """Return the nullable, randomly generated columns, in file order. diff --git a/src/dstrack/_track.py b/src/dstrack/_track.py index ed201f9..0c0c2c0 100644 --- a/src/dstrack/_track.py +++ b/src/dstrack/_track.py @@ -10,12 +10,13 @@ from dstrack import console from dstrack.errors import ( DatasetNotFoundError, + InputTooLargeError, StoreCorruptionError, StoreNotFoundError, ) from dstrack.paths import resolve_store_root from dstrack.readers import resolve_reader -from dstrack.snapshot import SnapshotBuilder +from dstrack.snapshot import SnapshotBuilder, StatsComputer from dstrack.store import write_snapshot @@ -72,6 +73,15 @@ def track( str | None, typer.Option("--created-by", help="Override the recorded author."), ] = None, + max_rows: Annotated[ + int | None, + typer.Option( + "--max-rows", + help="Maximum number of rows to process. Statistics are computed " + "in memory, so this bounds memory use; datasets larger than the " + "limit are rejected. Omit to use the default limit.", + ), + ] = None, ) -> None: """Compute a snapshot of a dataset and store it in the local store.""" # Locate the .dstrack/ store; without one there is nowhere to write. @@ -100,14 +110,23 @@ def track( # portable relative path is recorded; the real file on disk is what gets read # and hashed. console.info(f"Reading {path} and computing snapshot...") - snapshot = SnapshotBuilder().build( - tabular_reader, - dataset_name=name or path.stem, - dataset_path=dataset_path, - source=path, - source_type="file", # TODO: Add a way to manage other sources like folders - created_by=created_by or _default_created_by(), + # An explicit --max-rows overrides the StatsComputer default; omitting it + # keeps the built-in limit that guards against unbounded memory use. + stats_computer = ( + StatsComputer() if max_rows is None else StatsComputer(max_rows=max_rows) ) + try: + snapshot = SnapshotBuilder(stats_computer=stats_computer).build( + tabular_reader, + dataset_name=name or path.stem, + dataset_path=dataset_path, + source=path, + source_type="file", # TODO: Add a way to manage other sources like folders + created_by=created_by or _default_created_by(), + ) + except InputTooLargeError as e: + console.error(str(e)) + raise typer.Exit(code=1) from e # Persist the snapshot, attaching it to an existing lineage when one matches. try: diff --git a/src/dstrack/errors.py b/src/dstrack/errors.py index 02d942c..aadda49 100644 --- a/src/dstrack/errors.py +++ b/src/dstrack/errors.py @@ -12,3 +12,7 @@ class StoreCorruptionError(Exception): class DatasetNotFoundError(Exception): """Raised when a named dataset does not exist in the local store.""" + + +class InputTooLargeError(Exception): + """Raised when a dataset exceeds the configured snapshot row limit.""" diff --git a/src/dstrack/readers/_registry.py b/src/dstrack/readers/_registry.py index 18989b8..275684e 100644 --- a/src/dstrack/readers/_registry.py +++ b/src/dstrack/readers/_registry.py @@ -22,6 +22,7 @@ that *before* the class is ever instantiated. """ +import inspect import logging from collections.abc import Sequence from importlib.metadata import entry_points @@ -86,6 +87,13 @@ def check_reader_class(obj: object, *, origin: str) -> type[TabularReader]: "protocol: it must define a from_path(path) classmethod, which is " "how dstrack builds a reader it knows only by name." ) + # Check that from_path is a class method + if not isinstance(inspect.getattr_static(obj, "from_path"), classmethod): + raise TypeError( + f"{qualname} (from {origin}) defines from_path, but not as a " + "classmethod: dstrack calls it on the class (MyReader.from_path(path)), " + "so it must be decorated with @classmethod." + ) return obj @@ -101,12 +109,21 @@ def build_reader(reader_cls: type[TabularReader], path: Path) -> TabularReader: The reader instance. Raises: - TypeError: If ``reader_cls`` was never validated and lacks ``from_path``. + TypeError: If ``reader_cls`` was never validated and lacks ``from_path``, + or if its ``from_path`` returns something that is not a + [TabularReader][dstrack.readers._protocol.TabularReader]. """ factory: object = reader_cls if not isinstance(factory, ReaderFactory): raise TypeError(f"{reader_cls!r} is not an instance of {ReaderFactory}") - return factory.from_path(path) + reader: object = factory.from_path(path) + if not isinstance(reader, TabularReader): + qualname = f"{reader_cls.__module__}.{reader_cls.__qualname__}" + raise TypeError( + f"{qualname}.from_path returned {reader!r}, which is not a " + "TabularReader: from_path must build an object that follows this protocol." + ) + return reader def register_reader( @@ -133,6 +150,11 @@ def register_reader( never silently displaces an existing reader: two packages fighting over ``.parquet`` is a conflict the user has to see, not one to resolve by import order. + ValueError: If any extension is malformed. An extension is matched + against [Path.suffix][pathlib.PurePath.suffix], so it must be a + single leading-dot suffix such as ``".csv"``: a value like ``"csv"`` + or ``".tar.gz"`` could never match and is rejected rather than + silently registered as dead weight. """ check_reader_class(reader_cls, origin=f"reader {name!r}") @@ -146,6 +168,15 @@ def register_reader( ) claimed = [ext.lower() for ext in extensions] + for ext in claimed: + # An extension is only ever compared against Path.suffix, which is a + # single leading-dot component (".csv"). + if not ext.startswith(".") or Path(f"_{ext}").suffix != ext: + raise ValueError( + f"Extension {ext!r} (for reader {name!r}) is malformed: an " + "extension must be a single leading-dot suffix such as '.csv', " + "so that it can match Path.suffix." + ) for ext in claimed: if ext in _BY_EXTENSION: owner = _BY_EXTENSION[ext] diff --git a/src/dstrack/snapshot/_builder.py b/src/dstrack/snapshot/_builder.py index 4d9f987..20bd2e2 100644 --- a/src/dstrack/snapshot/_builder.py +++ b/src/dstrack/snapshot/_builder.py @@ -33,7 +33,9 @@ class SnapshotBuilder: [MetadataBuilder][dstrack.snapshot._metadata.MetadataBuilder]. stats_computer: Computer used for the per-column and dataset-level statistics. Defaults to a fresh - [StatsComputer][dstrack.snapshot._stats.StatsComputer]. + [StatsComputer][dstrack.snapshot._stats.StatsComputer]. Pass a + pre-configured ``StatsComputer(max_rows=...)`` here to change the + in-memory row limit the statistics pass enforces. """ def __init__( diff --git a/src/dstrack/snapshot/_stats.py b/src/dstrack/snapshot/_stats.py index b87b77f..451d5d3 100644 --- a/src/dstrack/snapshot/_stats.py +++ b/src/dstrack/snapshot/_stats.py @@ -3,6 +3,7 @@ from dataclasses import dataclass, field from typing import Protocol +from dstrack.errors import InputTooLargeError from dstrack.readers import Cell, TabularReader _NUMERIC_DTYPES: frozenset[str] = frozenset( @@ -10,6 +11,10 @@ ) _NUM_HISTOGRAM_BINS: int = 20 _TOP_VALUES_K: int = 50 +# Default ceiling on the number of rows a single snapshot pass will accumulate. +# The pass keeps per-column values, the distinct-row set, and per-string counts +# in memory, so worst-case memory grows with the row count; this bounds it. +_DEFAULT_MAX_ROWS: int = 10_000_000 @dataclass(frozen=True, slots=True) @@ -124,13 +129,24 @@ def stats(self, null_fraction: float) -> ColumnStats: ... @dataclass class _NumericAcc: - """Running accumulator for an ``int*``/``float*``/``bool`` column.""" + """Running accumulator for an ``int*``/``float*``/``bool`` column. + + Non-finite values (``NaN``, ``+inf``, ``-inf``) are treated as null: they + are counted in ``null_count`` and excluded from ``vals``. They carry no + usable magnitude and would otherwise poison mean/std/percentile/histogram + math (e.g. a ``NaN`` step makes the histogram bucket ``int(...)`` raise + ``ValueError``) and produce non-JSON-representable output. + """ vals: list[float] = field(default_factory=list) null_count: int = 0 def update(self, val: Cell) -> None: - self.vals.append(float(val)) # type: ignore[arg-type] + f = float(val) # type: ignore[arg-type] + if not math.isfinite(f): + self.null_count += 1 + return + self.vals.append(f) def is_constant(self) -> bool: return not self.vals or min(self.vals) == max(self.vals) @@ -227,8 +243,25 @@ class StatsComputer: Handles ``int*``, ``float*``, and ``bool`` dtypes as numeric, ``string`` columns, and ``datetime64`` columns. Unknown dtypes produce only null counts. + + The pass is exact but in-memory: per-column values, the distinct-row set, + and per-string value counts are all retained until it completes, so + worst-case memory grows with the row count. ``max_rows`` bounds that + growth by rejecting datasets larger than the limit rather than risking + memory exhaustion. + + Args: + max_rows: Maximum number of rows the pass will accept. When the reader + yields more than this, an + [InputTooLargeError][dstrack.errors.InputTooLargeError] is raised. + Defaults to + [_DEFAULT_MAX_ROWS][dstrack.snapshot._stats._DEFAULT_MAX_ROWS]. + Pass ``None`` to disable the limit when enough memory is available. """ + def __init__(self, *, max_rows: int | None = _DEFAULT_MAX_ROWS) -> None: + self._max_rows = max_rows + def compute(self, reader: TabularReader) -> DatasetStats: """Run a full data pass and return aggregated statistics. @@ -239,6 +272,9 @@ def compute(self, reader: TabularReader) -> DatasetStats: Returns: A populated [DatasetStats][dstrack.snapshot._stats.DatasetStats] instance. + + Raises: + InputTooLargeError: If the reader yields more than ``max_rows`` rows. """ cols = reader.columns() col_names = [c.name for c in cols] @@ -256,9 +292,22 @@ def compute(self, reader: TabularReader) -> DatasetStats: ) num_rows += batch_rows duplicate_count += batch_duplicates + # Checked per batch so an oversized file is rejected mid-read, + # before the whole dataset is pulled into memory. + self._check_row_limit(num_rows) return self._build_dataset_stats(col_names, accs, num_rows, duplicate_count) + def _check_row_limit(self, num_rows: int) -> None: + """Raise if ``num_rows`` has exceeded the configured ``max_rows``.""" + if self._max_rows is not None and num_rows > self._max_rows: + raise InputTooLargeError( + f"Dataset exceeds the snapshot row limit of " + f"{self._max_rows:,} rows. Statistics are computed in memory; " + f"raise the limit with --max-rows (or max_rows=None to disable) " + f"only if enough memory is available." + ) + def _init_accumulators(self, col_dtypes: dict[str, str]) -> dict[str, _ColumnAcc]: """Create one accumulator per column, chosen by dtype.""" accs: dict[str, _ColumnAcc] = {} diff --git a/src/dstrack/store.py b/src/dstrack/store.py index f028003..0639dfa 100644 --- a/src/dstrack/store.py +++ b/src/dstrack/store.py @@ -6,12 +6,24 @@ pointing at a snapshot that was not fully written. Re-tracking a source whose recorded path matches an existing dataset's latest snapshot continues that dataset's lineage rather than creating a new one. + +The three writes are not a single atomic transaction, so two safeguards keep +readers consistent. A store-wide advisory lock (:func:`_store_lock`) serializes +the whole resolve-parent-then-write sequence, so concurrent writers cannot read +the same ``HEAD`` as their parent nor mint two datasets for one path. And +``HEAD`` -- written last -- is the single source of truth for what is committed: +recovery and path matching read the log entry that ``HEAD`` names rather than +trusting the final line, so a ``log.jsonl`` left one entry ahead by a crash +between the append and the ``HEAD`` write is ignored. """ +import contextlib import json import os +import sys import tempfile import uuid +from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path from typing import Any, Final @@ -86,6 +98,8 @@ def write_snapshot( Raises: KeyError: If ``snapshot`` lacks ``snapshot_id`` or ``dataset_path``. + ValueError: If ``snapshot_id`` or an explicit ``dataset_id`` resolves to + a path outside the store. DatasetNotFoundError: If ``dataset_id`` is given but names no dataset in the store. StoreCorruptionError: If an existing dataset's ``log.jsonl`` ends in a @@ -96,36 +110,48 @@ def write_snapshot( dataset_path = snapshot["dataset_path"] snapshot_id = snapshot["snapshot_id"] - if dataset_id is not None: - # An explicit id continues a lineage; it must name a dataset that is - # already there, or a typo would mint a nameless one behind the user's - # back and report it as a continuation. - _check_dataset_exists(datasets_dir, dataset_id) - is_new_dataset = False - else: - dataset_id = _match_dataset_by_path(datasets_dir, dataset_path) - is_new_dataset = dataset_id is None - if dataset_id is None: - dataset_id = str(uuid.uuid4()) - - dataset_dir = datasets_dir / dataset_id - snapshots_dir = dataset_dir / "snapshots" - snapshots_dir.mkdir(parents=True, exist_ok=True) - - # HEAD is the dataset's latest snapshot, so it is this snapshot's parent. - # Absent for a dataset's first snapshot, whether it was just minted here or - # named by an explicit `dataset_id` that has no snapshots yet. - parent_snapshot_id = _read_head(dataset_dir) - payload = {**snapshot, "parent_snapshot_id": parent_snapshot_id} - - snapshot_path = snapshots_dir / f"{snapshot_id}.json" - _atomic_write(snapshot_path, json.dumps(payload, indent=2) + "\n") - - log_line = {key: payload.get(key) for key in _LOG_FIELDS} - with (dataset_dir / "log.jsonl").open("a", encoding="utf-8") as fh: - fh.write(json.dumps(log_line) + "\n") - - _atomic_write(dataset_dir / "HEAD", snapshot_id + "\n") + # Lock the store to avoid cross-process contamination + with _store_lock(store_root): + if dataset_id is not None: + # An explicit id continues a lineage; it must name a dataset + _ensure_direct_child( + base=datasets_dir, + candidate=datasets_dir / dataset_id, + kind="dataset_id", + value=dataset_id, + ) + _check_dataset_exists(datasets_dir, dataset_id) + is_new_dataset = False + else: + dataset_id = _match_dataset_by_path(datasets_dir, dataset_path) + is_new_dataset = dataset_id is None + if dataset_id is None: + dataset_id = str(uuid.uuid4()) + + dataset_dir = datasets_dir / dataset_id + snapshots_dir = dataset_dir / "snapshots" + snapshot_path = snapshots_dir / f"{snapshot_id}.json" + # snapshot_id arrives in the payload from outside, so confirm the file + # lands inside the dataset's snapshots/ before creating any directory. + _ensure_direct_child( + base=snapshots_dir, + candidate=snapshot_path, + kind="snapshot_id", + value=snapshot_id, + ) + snapshots_dir.mkdir(parents=True, exist_ok=True) + + # HEAD is the dataset's latest snapshot, so it is this snapshot's parent + parent_snapshot_id = _read_head(dataset_dir) + payload = {**snapshot, "parent_snapshot_id": parent_snapshot_id} + + _atomic_write(snapshot_path, json.dumps(payload, indent=2) + "\n") + + log_line = {key: payload.get(key) for key in _LOG_FIELDS} + with (dataset_dir / "log.jsonl").open("a", encoding="utf-8") as fh: + fh.write(json.dumps(log_line) + "\n") + + _atomic_write(dataset_dir / "HEAD", snapshot_id + "\n") return SnapshotWriteResult( dataset_id=dataset_id, @@ -136,6 +162,34 @@ def write_snapshot( ) +def _ensure_direct_child(base: Path, candidate: Path, kind: str, value: str) -> None: + """Raise unless ``candidate`` resolves to a direct child of ``base``. + + ``dataset_id`` and ``snapshot_id`` are joined onto the store to form a single + directory or file name directly under ``base``. A value holding a ``..`` + segment, a path separator, or an absolute path can make ``candidate`` land + outside ``base`` (and clobber files elsewhere) or nest under a subdirectory + that does not exist. Requiring the resolved parent to equal ``base`` rejects + both, regardless of how the escape is spelled, before any file is written. + + Args: + base: Directory the candidate must sit directly under, e.g. the store's + ``datasets/`` directory. + candidate: The joined path derived from ``value``. + kind: Name of the identifier, used in the error message. + value: The identifier value, used in the error message. + + Raises: + ValueError: If ``candidate`` does not resolve to a direct child of + ``base``. + """ + if candidate.resolve().parent != base.resolve(): + raise ValueError( + f"Invalid {kind} {value!r}: it must name a single entry directly " + "inside the store, not a nested or outside path." + ) + + def _check_dataset_exists(datasets_dir: Path, dataset_id: str) -> None: """Raise unless ``dataset_id`` names a dataset already in the store. @@ -183,47 +237,57 @@ def _match_dataset_by_path(datasets_dir: Path, dataset_path: str) -> str | None: for dataset_dir in sorted(datasets_dir.iterdir()): if not dataset_dir.is_dir(): continue - entry = _read_head_log_entry(dataset_dir) + entry = _read_committed_log_entry(dataset_dir) if entry is not None and entry.get("dataset_path") == dataset_path: return dataset_dir.name return None -def _read_head_log_entry(dataset_dir: Path) -> dict[str, Any] | None: - """Return the last (latest) parsed line of a dataset's ``log.jsonl``. +def _read_committed_log_entry(dataset_dir: Path) -> dict[str, Any] | None: + """Return the ``log.jsonl`` entry for the snapshot ``HEAD`` names. - The last line always corresponds to ``HEAD``: both are written by the same - [write_snapshot][dstrack.store.write_snapshot] call. + ``HEAD`` is written last, so it is the single source of truth for what is + committed. The final log line is not: a crash between the ``log.jsonl`` + append and the ``HEAD`` write can leave the log one entry ahead of ``HEAD``. + So the committed entry is the one whose ``snapshot_id`` equals ``HEAD``, + which is guaranteed to be present because its append precedes the ``HEAD`` + write that names it. Args: dataset_dir: A ``datasets//`` directory. Returns: - The parsed last non-blank line, or ``None`` if the dataset has no - ``log.jsonl`` or it holds no entries yet. + The parsed log entry the dataset's ``HEAD`` points at, or ``None`` if + the dataset has no ``HEAD`` or no ``log.jsonl`` yet. Raises: - StoreCorruptionError: If that last line is not valid JSON. + StoreCorruptionError: If a log line is not valid JSON, or no entry + matches ``HEAD`` (the log lost the committed snapshot's line). """ + head = _read_head(dataset_dir) + if head is None: + return None log_path = dataset_dir / "log.jsonl" if not log_path.is_file(): return None - last_line = "" with log_path.open(encoding="utf-8") as fh: for line in fh: - if line.strip(): - last_line = line - if not last_line: - return None - try: - entry: dict[str, Any] = json.loads(last_line) - except json.JSONDecodeError as e: - raise StoreCorruptionError( - f"The last entry of `{log_path}` is not valid JSON. The file may " - "have been truncated by an interrupted write; restore it from git " - "or delete the trailing line." - ) from e - return entry + if not line.strip(): + continue + try: + entry: dict[str, Any] = json.loads(line) + except json.JSONDecodeError as e: + raise StoreCorruptionError( + f"An entry of `{log_path}` is not valid JSON. The file may " + "have been truncated by an interrupted write; restore it " + "from git or delete the offending line." + ) from e + if entry.get("snapshot_id") == head: + return entry + raise StoreCorruptionError( + f"`{log_path}` has no entry for HEAD {head!r}. The log is missing the " + "committed snapshot's line; restore it from git." + ) def _read_head(dataset_dir: Path) -> str | None: @@ -268,3 +332,56 @@ def _atomic_write(path: Path, text: str) -> None: except BaseException: Path(tmp_name).unlink(missing_ok=True) raise + + +@contextlib.contextmanager +def _store_lock(store_root: Path) -> Iterator[None]: + """Hold an exclusive, cross-process lock on the store for the block's body. + + Serializes writers across processes so the resolve-parent-then-write + sequence in [write_snapshot][dstrack.store.write_snapshot] commits as a + unit: no two writers observe the same ``HEAD`` as a parent, and no two + path-matched writers mint separate datasets for one path. The lock is a + single ``.lock`` file at the store root, held for every dataset rather than + per-dataset, which also covers the cross-dataset scan that path matching + performs. It is an OS advisory lock, so it is released automatically if the + holding process dies, leaving no stale lock to clear by hand. + + Args: + store_root: Path to the ``.dstrack/`` directory. Created if absent so + the lock file has somewhere to live. + + Yields: + Nothing; the caller runs its critical section inside the ``with``. + """ + store_root.mkdir(parents=True, exist_ok=True) + # Open without truncating: the file is a lock handle, its contents unused. + fh = (store_root / ".lock").open("a", encoding="utf-8") + try: + _acquire(fh) + try: + yield + finally: + _release(fh) + finally: + fh.close() + + +if sys.platform == "win32": # pragma: no cover - platform specific + import msvcrt + + def _acquire(fh: Any) -> None: + msvcrt.locking(fh.fileno(), msvcrt.LK_LOCK, 1) + + def _release(fh: Any) -> None: + fh.seek(0) + msvcrt.locking(fh.fileno(), msvcrt.LK_UNLCK, 1) + +else: + import fcntl + + def _acquire(fh: Any) -> None: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX) + + def _release(fh: Any) -> None: + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py index 1ce5596..b3e27c1 100644 --- a/tests/test_benchmark_cli.py +++ b/tests/test_benchmark_cli.py @@ -3,6 +3,7 @@ import shutil from dataclasses import dataclass from pathlib import Path +from typing import ClassVar import pytest from typer.testing import CliRunner @@ -57,10 +58,11 @@ def _make_run() -> BenchmarkRun: class _StubRunner: """Stands in for BenchmarkRunner, recording how the CLI constructed it.""" + instances: ClassVar[list["_StubRunner"]] = [] + def __init__( self, spec: SyntheticCsvSpec, *, observer: object = None, profile: bool = True ) -> None: - self.instances: list[_StubRunner] = [] self.spec = spec self.observer = observer self.profile = profile diff --git a/tests/test_reader_registry.py b/tests/test_reader_registry.py index 21c2f2f..af531f4 100644 --- a/tests/test_reader_registry.py +++ b/tests/test_reader_registry.py @@ -76,6 +76,24 @@ def from_path(cls, path: Path) -> "NotAReader": return cls() +class InstanceFromPathReader: + """Reads fine and has from_path, but forgot the @classmethod decorator. + + Passes a bare ``isinstance(cls, ReaderFactory)`` check because the attribute + exists, yet ``cls.from_path(path)`` would bind ``path`` to ``self``. + """ + + def from_path(self, path: Path) -> "InstanceFromPathReader": + CONSTRUCTED.append(Path(path)) + return self + + def columns(self) -> list[ColumnInfo]: + return [] + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + return iter([]) + + not_a_class = "definitely not a class" @@ -264,6 +282,45 @@ def test_class_missing_from_path_is_rejected_unconstructed(tmp_path: Path) -> No assert CONSTRUCTED == [] +def test_instance_method_from_path_is_rejected_unconstructed(tmp_path: Path) -> None: + """A from_path that forgot @classmethod is caught before it is ever called.""" + path = tmp_path / "a.thing" + path.touch() + + with pytest.raises(TypeError, match="classmethod"): + resolve_reader(path, reader=spec_for(InstanceFromPathReader)) + + assert CONSTRUCTED == [] + + +def test_register_reader_rejects_instance_method_from_path() -> None: + """The non-classmethod check also guards registration, not just build time.""" + with pytest.raises(TypeError, match="classmethod"): + register_reader(InstanceFromPathReader, name="bad") # type: ignore[arg-type] + + assert "bad" not in available_readers() + + +def test_build_reader_rejects_non_reader_return(tmp_path: Path) -> None: + """build_reader rejects a from_path that returns something un-readable.""" + path = tmp_path / "a.thing" + path.touch() + + class BadReturnReader: + @classmethod + def from_path(cls, path: Path) -> object: + return "not a reader" + + def columns(self) -> list[ColumnInfo]: + return [] + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + return iter([]) + + with pytest.raises(TypeError, match="not a TabularReader"): + _registry.build_reader(BadReturnReader, path) + + def test_non_class_target_is_rejected() -> None: """A spec that names a module-level value, not a class, fails cleanly.""" with pytest.raises(TypeError, match="not a class"): @@ -295,6 +352,25 @@ def test_duplicate_extension_raises() -> None: register_reader(GoodReader, name="good", extensions=[".csv"]) +@pytest.mark.parametrize("ext", ["csv", "", ".", ".tar.gz"]) +def test_malformed_extension_raises(ext: str) -> None: + """An extension that could never match Path.suffix is rejected, not stored.""" + with pytest.raises(ValueError, match="malformed"): + register_reader(GoodReader, name="good", extensions=[ext]) + + assert ext not in known_extensions() + assert "good" not in available_readers() + + +def test_malformed_extension_leaves_registry_untouched() -> None: + """A bad extension is caught before the name or any extension is claimed.""" + with pytest.raises(ValueError, match="malformed"): + register_reader(GoodReader, name="good", extensions=[".fine", "bad"]) + + assert "good" not in available_readers() + assert ".fine" not in known_extensions() + + def test_explicit_extensions_override_the_class_attribute() -> None: """Passing extensions replaces EXTENSIONS rather than adding to it.""" register_reader(GoodReader, name="good", extensions=[".other"]) diff --git a/tests/test_snapshot_stats.py b/tests/test_snapshot_stats.py index 419de91..09fb675 100644 --- a/tests/test_snapshot_stats.py +++ b/tests/test_snapshot_stats.py @@ -1,10 +1,12 @@ """Tests for StatsComputer.""" +import math from collections.abc import Iterator from dataclasses import fields import pytest +from dstrack.errors import InputTooLargeError from dstrack.readers import Cell, ColumnInfo, TabularReader from dstrack.snapshot import DatasetStats, StatsComputer @@ -252,6 +254,62 @@ def test_bool_treated_as_numeric() -> None: assert s.max == 1.0 +def test_numeric_nan_treated_as_null() -> None: + """NaN values are counted as null and excluded from numeric statistics.""" + nan = float("nan") + stats = _compute( + [ColumnInfo("v", "float64")], + [[1.0], [nan], [3.0]], + ) + s = stats.column_stats["v"] + assert s.null_count == 1 + assert s.null_fraction == pytest.approx(1 / 3, abs=1e-9) + assert s.min == 1.0 + assert s.max == 3.0 + assert s.mean == pytest.approx(2.0, abs=1e-9) + assert s.num_unique == 2 + assert math.isfinite(s.std) + + +def test_numeric_infinity_treated_as_null() -> None: + """Positive and negative infinity are counted as null, keeping stats finite.""" + stats = _compute( + [ColumnInfo("v", "float64")], + [[1.0], [float("inf")], [float("-inf")], [5.0]], + ) + s = stats.column_stats["v"] + assert s.null_count == 2 + assert s.null_fraction == pytest.approx(0.5, abs=1e-9) + assert s.min == 1.0 + assert s.max == 5.0 + assert math.isfinite(s.mean) + assert all(math.isfinite(e) for e in s.histogram.bin_edges) + + +def test_numeric_all_non_finite_returns_zeros() -> None: + """A column of only non-finite values behaves like an all-null column.""" + stats = _compute( + [ColumnInfo("v", "float64")], + [[float("nan")], [float("inf")]], + ) + s = stats.column_stats["v"] + assert s.null_count == 2 + assert s.mean == 0.0 + assert s.std == 0.0 + assert s.num_unique == 0 + assert "v" in stats.constant_columns + + +def test_numeric_non_finite_histogram_does_not_raise() -> None: + """A NaN in the data must not break histogram bucketing (regression).""" + rows: list[list[Cell]] = [[float(i)] for i in range(10)] + rows.append([float("nan")]) + stats = _compute([ColumnInfo("v", "float64")], rows) + hist = stats.column_stats["v"].histogram + assert len(hist.bin_edges) == len(hist.counts) + 1 + assert sum(hist.counts) == 10 + + # --------------------------------------------------------------------------- # String column stats # --------------------------------------------------------------------------- @@ -362,3 +420,55 @@ def test_datetime_all_null() -> None: assert s.min == "" assert s.max == "" assert s.range_days == pytest.approx(0.0, abs=1e-9) + + +# --------------------------------------------------------------------------- +# max_rows limit +# --------------------------------------------------------------------------- + + +def _batched_reader( + cols: list[ColumnInfo], batches: list[list[list[Cell]]] +) -> TabularReader: + """A reader that yields pre-split batches, to exercise the per-batch guard.""" + + class _Stub: + def columns(self) -> list[ColumnInfo]: + return cols + + def iter_batches(self, batch_size: int = 1000) -> Iterator[list[list[Cell]]]: + yield from batches + + return _Stub() # type: ignore[return-value] + + +def test_max_rows_at_limit_computes_normally() -> None: + """A dataset exactly at the row limit is processed without error.""" + cols = [ColumnInfo("x", "int64")] + rows: list[list[Cell]] = [[i] for i in range(3)] + stats = StatsComputer(max_rows=3).compute(_reader(cols, rows)) + assert stats.num_rows == 3 + + +def test_max_rows_exceeded_raises() -> None: + """A dataset larger than the row limit raises InputTooLargeError.""" + cols = [ColumnInfo("x", "int64")] + rows: list[list[Cell]] = [[i] for i in range(5)] + with pytest.raises(InputTooLargeError, match="max-rows"): + StatsComputer(max_rows=3).compute(_reader(cols, rows)) + + +def test_max_rows_enforced_across_batches() -> None: + """The limit counts rows dataset-wide, not per batch.""" + cols = [ColumnInfo("x", "int64")] + batches: list[list[list[Cell]]] = [[[0], [1]], [[2], [3]]] + with pytest.raises(InputTooLargeError): + StatsComputer(max_rows=3).compute(_batched_reader(cols, batches)) + + +def test_max_rows_none_never_raises() -> None: + """max_rows=None disables the limit.""" + cols = [ColumnInfo("x", "int64")] + rows: list[list[Cell]] = [[i] for i in range(100)] + stats = StatsComputer(max_rows=None).compute(_reader(cols, rows)) + assert stats.num_rows == 100 diff --git a/tests/test_store.py b/tests/test_store.py index a706c9f..d467826 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -97,3 +97,48 @@ def test_log_line_appended_per_snapshot(tmp_path: Path) -> None: assert (tmp_path / "datasets" / first.dataset_id / "HEAD").read_text().strip() == ( second.snapshot_id ) + + +def test_trailing_log_entry_past_head_is_ignored(tmp_path: Path) -> None: + """Path matching follows committed HEAD, not a log line a crash left behind. + + A crash between the ``log.jsonl`` append and the ``HEAD`` write leaves the + log one entry ahead of what is committed. Re-tracking the path must continue + the lineage from HEAD and ignore that uncommitted trailing line. + """ + first = write_snapshot(_snapshot("data/ds.csv"), store_root=tmp_path) + dataset_dir = tmp_path / "datasets" / first.dataset_id + + # Simulate the crash: a log line for a snapshot HEAD never advanced onto, + # recording a different path than the committed HEAD. + phantom = {"snapshot_id": "phantom", "dataset_path": "data/other.csv"} + with (dataset_dir / "log.jsonl").open("a", encoding="utf-8") as fh: + fh.write(json.dumps(phantom) + "\n") + + second = write_snapshot(_snapshot("data/ds.csv"), store_root=tmp_path) + + assert not second.is_new_dataset + assert second.dataset_id == first.dataset_id + assert second.parent_snapshot_id == first.snapshot_id + + +@pytest.mark.parametrize("bad_id", ["../../evil", "a/b", "/etc/passwd"]) +def test_traversal_snapshot_id_is_rejected(tmp_path: Path, bad_id: str) -> None: + """A snapshot_id that would escape the dataset's snapshots/ is refused.""" + with pytest.raises(ValueError): + write_snapshot(_snapshot(snapshot_id=bad_id), store_root=tmp_path) + + +@pytest.mark.parametrize("bad_id", ["../../evil", "a/b", "/etc/passwd"]) +def test_traversal_dataset_id_is_rejected(tmp_path: Path, bad_id: str) -> None: + """An explicit dataset_id that would escape datasets/ is refused.""" + with pytest.raises(ValueError): + write_snapshot(_snapshot(), store_root=tmp_path, dataset_id=bad_id) + + +def test_rejected_snapshot_id_writes_nothing(tmp_path: Path) -> None: + """A rejected snapshot_id leaves no dataset directory behind.""" + with pytest.raises(ValueError): + write_snapshot(_snapshot(snapshot_id="../../evil"), store_root=tmp_path) + + assert not (tmp_path / "datasets").exists() diff --git a/tests/test_track.py b/tests/test_track.py index a849b9c..6ee7464 100644 --- a/tests/test_track.py +++ b/tests/test_track.py @@ -182,6 +182,30 @@ def test_track_unknown_dataset_id_errors(tmp_path: Path, store: Path) -> None: assert "does-not-exist" in result.output +def test_track_max_rows_below_row_count_exits_with_error( + tmp_path: Path, store: Path +) -> None: + """--max-rows below the dataset's row count fails with an actionable message.""" + csv = _write_csv(tmp_path / "data.csv", rows=5) + + result = runner.invoke(app, ["track", str(csv), "--max-rows", "2"]) + + assert result.exit_code == 1 + assert "--max-rows" in result.output + # Nothing is written when the limit is exceeded. + assert not list(store.glob("datasets/*/snapshots/*.json")) + + +def test_track_max_rows_above_row_count_succeeds(tmp_path: Path, store: Path) -> None: + """A --max-rows at or above the row count writes the snapshot normally.""" + csv = _write_csv(tmp_path / "data.csv", rows=3) + + result = runner.invoke(app, ["track", str(csv), "--max-rows", "3"]) + + assert result.exit_code == 0, result.output + assert _load_only_snapshot(store)["num_rows"] == 3 + + # --------------------------------------------------------------------------- # Failure paths # --------------------------------------------------------------------------- From 0ea9391599f006db02cf20ef34007865977912ae Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 18:15:23 +0200 Subject: [PATCH 10/11] Fixes test failure on python 3.14 due to pathlib suffix interpretation. --- src/dstrack/readers/_registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/dstrack/readers/_registry.py b/src/dstrack/readers/_registry.py index 275684e..5339eba 100644 --- a/src/dstrack/readers/_registry.py +++ b/src/dstrack/readers/_registry.py @@ -171,7 +171,7 @@ def register_reader( for ext in claimed: # An extension is only ever compared against Path.suffix, which is a # single leading-dot component (".csv"). - if not ext.startswith(".") or Path(f"_{ext}").suffix != ext: + if not ext.startswith(".") or ext == "." or Path(f"_{ext}").suffix != ext: raise ValueError( f"Extension {ext!r} (for reader {name!r}) is malformed: an " "extension must be a single leading-dot suffix such as '.csv', " From 72021d1697dae47ca50cc2ac5204297183fc9052 Mon Sep 17 00:00:00 2001 From: Leonardo Ayala Date: Wed, 15 Jul 2026 18:31:04 +0200 Subject: [PATCH 11/11] fixes test dataclass initialization to use field with default factory instead of overwriting in __init__ --- tests/test_benchmark_cli.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py index b3e27c1..2bc115d 100644 --- a/tests/test_benchmark_cli.py +++ b/tests/test_benchmark_cli.py @@ -1,7 +1,7 @@ """Tests for the benchmark CLI in src/dstrack/_benchmark/_cli.py.""" import shutil -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import ClassVar @@ -24,12 +24,8 @@ class _RecordingReporter: """Stands in for ConsoleReporter, recording what the CLI asked it to print.""" - kept_paths: list[Path] - reported: list[tuple[BenchmarkRun, int]] - - def __init__(self) -> None: - self.kept_paths = [] - self.reported = [] + kept_paths: list[Path] = field(default_factory=list) + reported: list[tuple[BenchmarkRun, int]] = field(default_factory=list) def report(self, run: BenchmarkRun, *, profile_limit: int) -> None: self.reported.append((run, profile_limit))