diff --git a/conformance/text.json b/conformance/text.json index c162237..1f32d01 100644 --- a/conformance/text.json +++ b/conformance/text.json @@ -21,6 +21,7 @@ { "name": "one BOM is removed from the start of every record", "input": "\ufeffa\n\ufeffb", "options": { "trim": false }, "expect": ["a", "b"] }, { "name": "only one BOM per record", "input": "\ufeff\ufeffa", "options": { "trim": false }, "expect": ["\ufeffa"] }, { "name": "trim uses the ECMAScript whitespace set", "input": "\u00a0a\u3000\n\u2028b\ufeff\n\u205fc\u1680", "expect": ["a", "b", "c"] }, + { "name": "a row of only Unicode whitespace is blank", "input": "a\n\u00a0\u3000\ufeff\nb", "expect": ["a", "b"] }, { "name": "NEL is not whitespace to trim", "input": "\u0085a\u0085", "expect": ["\u0085a\u0085"] }, { "name": "information separators are not whitespace to trim", "input": "\u001ca\u001f", "expect": ["\u001ca\u001f"] }, { "name": "invalid UTF-8 decodes lossily, one replacement per maximal subpart", "input": { "base64": "Yf9iCuKC" }, "expect": ["a\ufffdb", "\ufffd"] }, diff --git a/crates/spliterator-py/README.md b/crates/spliterator-py/README.md index 7956129..8d20bb5 100644 --- a/crates/spliterator-py/README.md +++ b/crates/spliterator-py/README.md @@ -26,6 +26,18 @@ for line in spliterator.read_text(b" a \n\n b\n"): # "a", "b" ... ``` +## Async and counts + +Each reader has an async twin that reads each batch on a worker thread, so the event loop never blocks on I/O: + +```python +async for row in spliterator.aread_csv("people.csv"): + ... +``` + +`count_text`, `count_csv`, `count_tsv` and `count_psv` return the number of rows the matching reader would yield, +without building them. + ## Sources A source is a path (`str` or `os.PathLike`), bytes-like data, or a binary file object. diff --git a/crates/spliterator-py/python/spliterator/__init__.py b/crates/spliterator-py/python/spliterator/__init__.py index df951c1..eb3f29e 100644 --- a/crates/spliterator-py/python/spliterator/__init__.py +++ b/crates/spliterator-py/python/spliterator/__init__.py @@ -8,6 +8,9 @@ ``path.gz`` in its place when one exists (``prefer_compressed``), and any source beginning with the zstd or gzip magic bytes is decompressed before it is split (``decompress``). Nothing is opened until the first row is pulled. +Every reader has an async twin (``aread_text``, ``aread_csv``, …) that reads each batch of rows on a worker thread, and +text and the CSV family have counts (``count_text``, ``count_csv``, …) that equal the number of rows the reader yields. + >>> import spliterator >>> list(spliterator.read_csv(b"name,age\\nAda,36\\n")) [{'name': 'Ada', 'age': '36'}] @@ -15,21 +18,34 @@ from __future__ import annotations +import asyncio import json import os -from collections.abc import Callable, Iterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping from importlib.metadata import PackageNotFoundError, version -from typing import Any, BinaryIO, Literal, TypeAlias +from typing import Any, BinaryIO, Literal, TypeAlias, TypedDict, TypeVar, Unpack from . import _native from ._errors import DuplicateHeaderError, JSONRowError, SpliteratorError, UnterminatedQuoteError __all__ = [ + "CSVOptions", "DuplicateHeaderError", + "JSONLOptions", "JSONRowError", "Source", "SpliteratorError", + "TextOptions", "UnterminatedQuoteError", + "aread_csv", + "aread_jsonl", + "aread_psv", + "aread_text", + "aread_tsv", + "count_csv", + "count_psv", + "count_text", + "count_tsv", "read_csv", "read_jsonl", "read_psv", @@ -49,6 +65,147 @@ Delimiter: TypeAlias = str | bytes +Transformers: TypeAlias = Mapping[str, Callable[[str], Any]] + +T = TypeVar("T") + + +class TextOptions(TypedDict, total=False): + """Options for :func:`read_text`, :func:`aread_text` and :func:`count_text`.""" + + delimiter: Delimiter + """The record delimiter. Default ``"\\n"``.""" + trim: bool + """Strip the ECMAScript whitespace set (not ``str.strip``'s) from each row. Default ``True``.""" + skip_empty: bool + """Drop empty records, and with ``trim`` rows that trim to nothing. Default ``True``.""" + crlf: bool + """Treat one ``\\r`` before a delimiter as part of it. Default ``False``; ``trim`` removes it anyway.""" + drop: int + """Rows to skip, counted after skipping. Default ``0``.""" + take: int | None + """Rows to yield at most, counted after skipping. Default unlimited.""" + enable_quote_handling: bool + """Keep a delimiter inside ``"…"`` from splitting. Default ``False``.""" + allow_unterminated_quote: bool + """Read input ending inside quotes instead of raising :class:`UnterminatedQuoteError`. Default ``False``.""" + decompress: bool + """Decompress a source beginning with the zstd or gzip magic bytes. Default ``True``.""" + prefer_compressed: bool + """Read ``path.zst`` or ``path.gz`` in place of ``path`` when one exists. Default ``True``.""" + + +class JSONLOptions(TypedDict, total=False): + """Options for :func:`read_jsonl` and :func:`aread_jsonl`.""" + + delimiter: Delimiter + """The record delimiter. Default ``"\\n"``.""" + skip_empty: bool + """Skip empty and whitespace-only rows. Default ``True``.""" + crlf: bool + """Treat one ``\\r`` before a delimiter as part of it. Default ``False``.""" + comment: str | list[str] | None + """Skip rows whose first non-whitespace text starts with one of these prefixes. Default none.""" + drop: int + """Rows to skip, counted after skipping. Default ``0``.""" + take: int | None + """Rows to parse at most, counted after skipping. Default unlimited.""" + decompress: bool + """Decompress a source beginning with the zstd or gzip magic bytes. Default ``True``.""" + prefer_compressed: bool + """Read ``path.zst`` or ``path.gz`` in place of ``path`` when one exists. Default ``True``.""" + + +class CSVOptions(TypedDict, total=False): + """Options for the CSV, TSV and PSV readers and counts. The defaults shown are CSV's.""" + + column_delimiter: Delimiter + """The column delimiter: ``","`` for CSV, ``"\\t"`` for TSV, ``"|"`` for PSV.""" + delimiter: Delimiter + """The row delimiter. Default ``"\\n"``.""" + header: bool + """Take the first record as the header. Default ``True``.""" + mode: Mode | None + """``"object"`` (default with a header), ``"array"`` (default without), or ``"entries"``.""" + normalize_keys: bool + """Snake-case header names and suffix repeats. Default ``False``.""" + enable_quote_handling: bool + """Keep delimiters inside ``"…"`` from splitting and unquote cells. ``True`` for CSV and PSV, ``False`` for TSV.""" + allow_unterminated_quote: bool + """Read input ending inside quotes instead of raising :class:`UnterminatedQuoteError`. Default ``False``.""" + crlf: bool + """Treat one ``\\r`` before a row delimiter as part of it. Default ``True``.""" + trim: bool + """Trim each cell, header included, after unquoting. Default ``True``.""" + skip_empty: bool + """Drop records whose bytes are empty after CRLF removal. Default ``True``.""" + drop: int + """Data rows to skip. Default ``0``.""" + take: int | None + """Data rows to yield at most. Default unlimited.""" + transformers: Transformers | None + """Functions applied by header name to object and entries rows, after trimming. Default none.""" + decompress: bool + """Decompress a source beginning with the zstd or gzip magic bytes. Default ``True``.""" + prefer_compressed: bool + """Read ``path.zst`` or ``path.gz`` in place of ``path`` when one exists. Default ``True``.""" + + +_TEXT_DEFAULTS: TextOptions = { + "delimiter": "\n", + "trim": True, + "skip_empty": True, + "crlf": False, + "drop": 0, + "take": None, + "enable_quote_handling": False, + "allow_unterminated_quote": False, + "decompress": True, + "prefer_compressed": True, +} + +_JSONL_DEFAULTS: JSONLOptions = { + "delimiter": "\n", + "skip_empty": True, + "crlf": False, + "comment": None, + "drop": 0, + "take": None, + "decompress": True, + "prefer_compressed": True, +} + +_CSV_DEFAULTS: CSVOptions = { + "column_delimiter": ",", + "delimiter": "\n", + "header": True, + "mode": None, + "normalize_keys": False, + "enable_quote_handling": True, + "allow_unterminated_quote": False, + "crlf": True, + "trim": True, + "skip_empty": True, + "drop": 0, + "take": None, + "transformers": None, + "decompress": True, + "prefer_compressed": True, +} + +_TSV_DEFAULTS: CSVOptions = {**_CSV_DEFAULTS, "column_delimiter": "\t", "enable_quote_handling": False} + +_PSV_DEFAULTS: CSVOptions = {**_CSV_DEFAULTS, "column_delimiter": "|"} + + +def _merge(defaults: Mapping[str, Any], options: Mapping[str, Any], reader: str) -> dict[str, Any]: + unknown = options.keys() - defaults.keys() + + if unknown: + raise TypeError(f"{reader}() got unexpected keyword arguments: {', '.join(sorted(unknown))}") + + return {**defaults, **options} + def _source(source: Source) -> str | bytes | BinaryIO: if isinstance(source, (bytes, bytearray, memoryview)): @@ -64,42 +221,76 @@ def _delimiter(value: Delimiter) -> bytes: return value.encode("utf-8") if isinstance(value, str) else bytes(value) -def read_text( - source: Source, - *, - delimiter: Delimiter = "\n", - trim: bool = True, - skip_empty: bool = True, - crlf: bool = False, - drop: int = 0, - take: int | None = None, - enable_quote_handling: bool = False, - allow_unterminated_quote: bool = False, - decompress: bool = True, - prefer_compressed: bool = True, -) -> Iterator[str]: +def _flatten(batches: Iterable[list[T]]) -> Iterator[T]: + for batch in batches: + yield from batch + + +async def _aflatten(open_batches: Callable[[], Iterator[list[T]]]) -> AsyncIterator[T]: + # Opening a path and reading a batch both block, so both happen on a worker thread; the rows are handed out here. + batches = await asyncio.to_thread(open_batches) + + while (batch := await asyncio.to_thread(next, batches, None)) is not None: + for row in batch: + yield row + + +# --- text ------------------------------------------------------------------------------------------------------------- + + +def _text_batches(source: Source, options: Mapping[str, Any]) -> _native.TextBatches: + o = _merge(_TEXT_DEFAULTS, options, "read_text") + + return _native.TextBatches( + _source(source), + _delimiter(o["delimiter"]), + o["trim"], + o["skip_empty"], + o["crlf"], + o["drop"], + o["take"], + o["enable_quote_handling"], + o["allow_unterminated_quote"], + o["decompress"], + o["prefer_compressed"], + ) + + +def read_text(source: Source, **options: Unpack[TextOptions]) -> Iterator[str]: """Yield the rows of a delimited text source. Each record is decoded as UTF-8 (each invalid sequence becomes U+FFFD) with one leading BOM removed. ``trim`` (on by default) strips the ECMAScript whitespace set, not ``str.strip``'s, and with ``skip_empty`` a row that trims to nothing is dropped. ``drop`` and ``take`` count the rows yielded. """ - batches = _native.TextBatches( + yield from _flatten(_text_batches(source, options)) + + +def aread_text(source: Source, **options: Unpack[TextOptions]) -> AsyncIterator[str]: + """:func:`read_text`, reading each batch of rows on a worker thread.""" + return _aflatten(lambda: _text_batches(source, options)) + + +def count_text(source: Source, **options: Unpack[TextOptions]) -> int: + """The number of rows :func:`read_text` would yield, without building them.""" + o = _merge(_TEXT_DEFAULTS, options, "count_text") + + return _native.count_text( _source(source), - _delimiter(delimiter), - trim, - skip_empty, - crlf, - drop, - take, - enable_quote_handling, - allow_unterminated_quote, - decompress, - prefer_compressed, + _delimiter(o["delimiter"]), + o["trim"], + o["skip_empty"], + o["crlf"], + o["drop"], + o["take"], + o["enable_quote_handling"], + o["allow_unterminated_quote"], + o["decompress"], + o["prefer_compressed"], ) - for batch in batches: - yield from batch + +# --- JSONL ------------------------------------------------------------------------------------------------------------ def _reject_constant(name: str) -> Any: @@ -110,133 +301,153 @@ def _reject_constant(name: str) -> Any: _JSON_DECODER = json.JSONDecoder(parse_constant=_reject_constant) -def read_jsonl( - source: Source, - *, - delimiter: Delimiter = "\n", - skip_empty: bool = True, - crlf: bool = False, - comment: str | list[str] | None = None, - drop: int = 0, - take: int | None = None, - decompress: bool = True, - prefer_compressed: bool = True, -) -> Iterator[Any]: +def _jsonl_batches(source: Source, options: Mapping[str, Any]) -> _native.JsonlRows: + o = _merge(_JSONL_DEFAULTS, options, "read_jsonl") + comment = o["comment"] + + return _native.JsonlRows( + _source(source), + _delimiter(o["delimiter"]), + o["skip_empty"], + o["crlf"], + [comment] if isinstance(comment, str) else comment, + o["drop"], + o["take"], + o["decompress"], + o["prefer_compressed"], + ) + + +class _JSONParser: + """Parses rows in order, numbering them for the error a malformed row raises.""" + + def __init__(self) -> None: + self.row = 0 + + def __call__(self, text: str) -> Any: + try: + value = _JSON_DECODER.decode(text) + except ValueError as error: + raise JSONRowError(f"Failed to parse JSON at row {self.row}: {error}", self.row) from error + + self.row += 1 + + return value + + +def read_jsonl(source: Source, **options: Unpack[JSONLOptions]) -> Iterator[Any]: """Yield each row of a JSONL source, parsed. Empty and whitespace-only rows are skipped. ``comment`` prefixes skip rows whose first non-whitespace text starts with one. ``drop`` and ``take`` count the rows that reach the parser. A row that fails to parse raises :class:`JSONRowError` naming its index among those rows. """ - comments = [comment] if isinstance(comment, str) else comment - batches = _native.JsonlRows( - _source(source), - _delimiter(delimiter), - skip_empty, - crlf, - comments, - drop, - take, - decompress, - prefer_compressed, - ) - row = 0 + parse = _JSONParser() - for batch in batches: - for text in batch: - try: - value = _JSON_DECODER.decode(text) - except ValueError as error: - raise JSONRowError(f"Failed to parse JSON at row {row}: {error}", row) from error + for text in _flatten(_jsonl_batches(source, options)): + yield parse(text) - yield value - row += 1 +async def aread_jsonl(source: Source, **options: Unpack[JSONLOptions]) -> AsyncIterator[Any]: + """:func:`read_jsonl`, reading each batch of rows on a worker thread. Parsing happens on the event loop.""" + parse = _JSONParser() -Transformers: TypeAlias = Mapping[str, Callable[[str], Any]] + async for text in _aflatten(lambda: _jsonl_batches(source, options)): + yield parse(text) + + +# --- CSV, TSV, PSV ---------------------------------------------------------------------------------------------------- -def _read_delimited( - source: Source, - *, - column_delimiter: Delimiter, - delimiter: Delimiter, - header: bool, - mode: Mode | None, - normalize_keys: bool, - enable_quote_handling: bool, - allow_unterminated_quote: bool, - crlf: bool, - trim: bool, - skip_empty: bool, - drop: int, - take: int | None, - transformers: Transformers | None, - decompress: bool, - prefer_compressed: bool, -) -> Iterator[Any]: - batches = _native.CsvBatches( +def _csv_batches(source: Source, o: Mapping[str, Any]) -> _native.CsvBatches: + return _native.CsvBatches( _source(source), - _delimiter(delimiter), - _delimiter(column_delimiter), - enable_quote_handling, - allow_unterminated_quote, - crlf, - trim, - skip_empty, - header, - mode, - normalize_keys, - drop, - take, - decompress, - prefer_compressed, + _delimiter(o["delimiter"]), + _delimiter(o["column_delimiter"]), + o["enable_quote_handling"], + o["allow_unterminated_quote"], + o["crlf"], + o["trim"], + o["skip_empty"], + o["header"], + o["mode"], + o["normalize_keys"], + o["drop"], + o["take"], + o["decompress"], + o["prefer_compressed"], ) - resolved_mode = mode or ("object" if header else "array") - # Transformers bind to header names and apply to object and entries rows, as `CSVSpliterator`'s do. - if not transformers or resolved_mode == "array": - for batch in batches: - yield from batch - return +class _Transformer: + """Applies transformers to object and entries rows, by header name, as `CSVSpliterator` does.""" - bound: dict[str, Callable[[str], Any]] | None = None + def __init__(self, batches: _native.CsvBatches, o: Mapping[str, Any]) -> None: + self.batches = batches + self.mode = o["mode"] or ("object" if o["header"] else "array") + self.transformers: Transformers | None = o["transformers"] if self.mode != "array" else None + self.bound: dict[str, Callable[[str], Any]] | None = None - for batch in batches: - if bound is None: - headers = set(batches.headers or ()) - bound = {name: transform for name, transform in transformers.items() if name in headers} + def batch(self, rows: list[Any]) -> list[Any]: + if not self.transformers: + return rows - for row in batch: - if resolved_mode == "object": + if self.bound is None: + headers = set(self.batches.headers or ()) + self.bound = {name: fn for name, fn in self.transformers.items() if name in headers} + + bound = self.bound + + if self.mode == "object": + for row in rows: for name, transform in bound.items(): row[name] = transform(row[name]) - else: - row = [(key, bound[key](value) if key in bound else value, index) for key, value, index in row] - yield row + return rows + + return [ + [(key, bound[key](value) if key in bound else value, index) for key, value, index in row] for row in rows + ] + + +def _transformed(source: Source, o: Mapping[str, Any]) -> Iterator[list[Any]]: + batches = _csv_batches(source, o) + transformer = _Transformer(batches, o) + + for batch in batches: + yield transformer.batch(batch) + + +def _read_csv(source: Source, defaults: CSVOptions, options: Mapping[str, Any], reader: str) -> Iterator[Any]: + yield from _flatten(_transformed(source, _merge(defaults, options, reader))) + + +def _aread_csv(source: Source, defaults: CSVOptions, options: Mapping[str, Any], reader: str) -> AsyncIterator[Any]: + o = _merge(defaults, options, reader) + + return _aflatten(lambda: _transformed(source, o)) -def read_csv( - source: Source, - *, - column_delimiter: Delimiter = ",", - delimiter: Delimiter = "\n", - header: bool = True, - mode: Mode | None = None, - normalize_keys: bool = False, - enable_quote_handling: bool = True, - allow_unterminated_quote: bool = False, - crlf: bool = True, - trim: bool = True, - skip_empty: bool = True, - drop: int = 0, - take: int | None = None, - transformers: Transformers | None = None, - decompress: bool = True, - prefer_compressed: bool = True, -) -> Iterator[Any]: +def _count_csv(source: Source, defaults: CSVOptions, options: Mapping[str, Any], reader: str) -> int: + o = _merge(defaults, options, reader) + + return _native.count_csv( + _source(source), + _delimiter(o["delimiter"]), + _delimiter(o["column_delimiter"]), + o["enable_quote_handling"], + o["allow_unterminated_quote"], + o["crlf"], + o["skip_empty"], + o["header"], + o["drop"], + o["take"], + o["decompress"], + o["prefer_compressed"], + ) + + +def read_csv(source: Source, **options: Unpack[CSVOptions]) -> Iterator[Any]: """Yield the rows of a CSV source. With a header (the default) rows are ``dict`` keyed by the header as written; without one they are ``list[str]``. @@ -247,102 +458,46 @@ def read_csv( a repeated header raises :class:`DuplicateHeaderError` unless ``normalize_keys`` suffixes it. Input ending inside quotes raises :class:`UnterminatedQuoteError` after the rows before it. ``drop`` and ``take`` count data rows. """ - return _read_delimited( - source, - column_delimiter=column_delimiter, - delimiter=delimiter, - header=header, - mode=mode, - normalize_keys=normalize_keys, - enable_quote_handling=enable_quote_handling, - allow_unterminated_quote=allow_unterminated_quote, - crlf=crlf, - trim=trim, - skip_empty=skip_empty, - drop=drop, - take=take, - transformers=transformers, - decompress=decompress, - prefer_compressed=prefer_compressed, - ) + return _read_csv(source, _CSV_DEFAULTS, options, "read_csv") -def read_tsv( - source: Source, - *, - column_delimiter: Delimiter = "\t", - delimiter: Delimiter = "\n", - header: bool = True, - mode: Mode | None = None, - normalize_keys: bool = False, - enable_quote_handling: bool = False, - allow_unterminated_quote: bool = False, - crlf: bool = True, - trim: bool = True, - skip_empty: bool = True, - drop: int = 0, - take: int | None = None, - transformers: Transformers | None = None, - decompress: bool = True, - prefer_compressed: bool = True, -) -> Iterator[Any]: +def read_tsv(source: Source, **options: Unpack[CSVOptions]) -> Iterator[Any]: """Yield the rows of a TSV source. As :func:`read_csv`, except that the TSV format has no quoting, so a ``"`` is data unless ``enable_quote_handling`` is set.""" - return _read_delimited( - source, - column_delimiter=column_delimiter, - delimiter=delimiter, - header=header, - mode=mode, - normalize_keys=normalize_keys, - enable_quote_handling=enable_quote_handling, - allow_unterminated_quote=allow_unterminated_quote, - crlf=crlf, - trim=trim, - skip_empty=skip_empty, - drop=drop, - take=take, - transformers=transformers, - decompress=decompress, - prefer_compressed=prefer_compressed, - ) + return _read_csv(source, _TSV_DEFAULTS, options, "read_tsv") -def read_psv( - source: Source, - *, - column_delimiter: Delimiter = "|", - delimiter: Delimiter = "\n", - header: bool = True, - mode: Mode | None = None, - normalize_keys: bool = False, - enable_quote_handling: bool = True, - allow_unterminated_quote: bool = False, - crlf: bool = True, - trim: bool = True, - skip_empty: bool = True, - drop: int = 0, - take: int | None = None, - transformers: Transformers | None = None, - decompress: bool = True, - prefer_compressed: bool = True, -) -> Iterator[Any]: +def read_psv(source: Source, **options: Unpack[CSVOptions]) -> Iterator[Any]: """Yield the rows of a pipe-separated source. As :func:`read_csv`, split on ``|``.""" - return _read_delimited( - source, - column_delimiter=column_delimiter, - delimiter=delimiter, - header=header, - mode=mode, - normalize_keys=normalize_keys, - enable_quote_handling=enable_quote_handling, - allow_unterminated_quote=allow_unterminated_quote, - crlf=crlf, - trim=trim, - skip_empty=skip_empty, - drop=drop, - take=take, - transformers=transformers, - decompress=decompress, - prefer_compressed=prefer_compressed, - ) + return _read_csv(source, _PSV_DEFAULTS, options, "read_psv") + + +def aread_csv(source: Source, **options: Unpack[CSVOptions]) -> AsyncIterator[Any]: + """:func:`read_csv`, reading each batch of rows on a worker thread.""" + return _aread_csv(source, _CSV_DEFAULTS, options, "aread_csv") + + +def aread_tsv(source: Source, **options: Unpack[CSVOptions]) -> AsyncIterator[Any]: + """:func:`read_tsv`, reading each batch of rows on a worker thread.""" + return _aread_csv(source, _TSV_DEFAULTS, options, "aread_tsv") + + +def aread_psv(source: Source, **options: Unpack[CSVOptions]) -> AsyncIterator[Any]: + """:func:`read_psv`, reading each batch of rows on a worker thread.""" + return _aread_csv(source, _PSV_DEFAULTS, options, "aread_psv") + + +def count_csv(source: Source, **options: Unpack[CSVOptions]) -> int: + """The number of data rows :func:`read_csv` would yield, without splitting cells. The header is skipped but not + read, so a repeated header is not an error here; ``mode``, ``normalize_keys`` and ``transformers`` are ignored.""" + return _count_csv(source, _CSV_DEFAULTS, options, "count_csv") + + +def count_tsv(source: Source, **options: Unpack[CSVOptions]) -> int: + """The number of data rows :func:`read_tsv` would yield.""" + return _count_csv(source, _TSV_DEFAULTS, options, "count_tsv") + + +def count_psv(source: Source, **options: Unpack[CSVOptions]) -> int: + """The number of data rows :func:`read_psv` would yield.""" + return _count_csv(source, _PSV_DEFAULTS, options, "count_psv") diff --git a/crates/spliterator-py/python/spliterator/_native.pyi b/crates/spliterator-py/python/spliterator/_native.pyi index b3323c6..fc2cb08 100644 --- a/crates/spliterator-py/python/spliterator/_native.pyi +++ b/crates/spliterator-py/python/spliterator/_native.pyi @@ -57,3 +57,31 @@ class CsvBatches(Iterator[list[Any]]): @property def headers(self) -> list[str] | None: ... def __next__(self) -> list[Any]: ... + +def count_text( + source: _Source, + delimiter: bytes, + trim: bool, + skip_empty: bool, + crlf: bool, + drop: int, + take: int | None, + enable_quote_handling: bool, + allow_unterminated_quote: bool, + decompress: bool, + prefer_compressed: bool, +) -> int: ... +def count_csv( + source: _Source, + row_delimiter: bytes, + column_delimiter: bytes, + enable_quote_handling: bool, + allow_unterminated_quote: bool, + crlf: bool, + skip_empty: bool, + header: bool, + drop: int, + take: int | None, + decompress: bool, + prefer_compressed: bool, +) -> int: ... diff --git a/crates/spliterator-py/src/lib.rs b/crates/spliterator-py/src/lib.rs index 6cd0ada..10dde7a 100644 --- a/crates/spliterator-py/src/lib.rs +++ b/crates/spliterator-py/src/lib.rs @@ -493,8 +493,90 @@ impl CsvBatches { } } +/// The number of text rows `TextBatches` would yield, counted with the interpreter released. +#[pyfunction] +#[pyo3(signature = ( + source, delimiter, trim, skip_empty, crlf, drop, take, enable_quote_handling, allow_unterminated_quote, + decompress, prefer_compressed +))] +#[allow(clippy::too_many_arguments)] +fn count_text( + py: Python<'_>, + source: &Bound<'_, PyAny>, + delimiter: Vec, + trim: bool, + skip_empty: bool, + crlf: bool, + drop: usize, + take: Option, + enable_quote_handling: bool, + allow_unterminated_quote: bool, + decompress: bool, + prefer_compressed: bool, +) -> PyResult { + let options = text::Options { + delimiter, + trim, + skip_empty, + crlf, + drop, + take, + enable_quote_handling, + allow_unterminated_quote, + open: open_options(decompress, prefer_compressed), + }; + let source = to_source(source)?; + + py.detach(|| text::count(source, options)) + .map_err(|error| to_py_error(py, error)) +} + +/// The number of data rows `CsvBatches` would yield, counted with the interpreter released. The header is skipped, +/// not read. +#[pyfunction] +#[pyo3(signature = ( + source, row_delimiter, column_delimiter, enable_quote_handling, allow_unterminated_quote, crlf, skip_empty, + header, drop, take, decompress, prefer_compressed +))] +#[allow(clippy::too_many_arguments)] +fn count_csv( + py: Python<'_>, + source: &Bound<'_, PyAny>, + row_delimiter: Vec, + column_delimiter: Vec, + enable_quote_handling: bool, + allow_unterminated_quote: bool, + crlf: bool, + skip_empty: bool, + header: bool, + drop: usize, + take: Option, + decompress: bool, + prefer_compressed: bool, +) -> PyResult { + let options = csv::Options { + row_delimiter, + column_delimiter, + enable_quote_handling, + allow_unterminated_quote, + crlf, + skip_empty, + header, + drop, + take, + open: open_options(decompress, prefer_compressed), + ..csv::Options::csv() + }; + let source = to_source(source)?; + + py.detach(|| csv::count(source, options)) + .map_err(|error| to_py_error(py, error)) +} + #[pymodule] fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(count_text, module)?)?; + module.add_function(wrap_pyfunction!(count_csv, module)?)?; module.add_class::()?; module.add_class::()?; module.add_class::()?; diff --git a/crates/spliterator-py/tests/test_conformance.py b/crates/spliterator-py/tests/test_conformance.py index f7240d2..da258ff 100644 --- a/crates/spliterator-py/tests/test_conformance.py +++ b/crates/spliterator-py/tests/test_conformance.py @@ -6,10 +6,11 @@ from __future__ import annotations +import asyncio import base64 import io import json -from collections.abc import Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from pathlib import Path from typing import Any @@ -44,6 +45,21 @@ "columnDelimiter": "column_delimiter", } +ASYNC_READERS: dict[str, Callable[..., AsyncIterator[Any]]] = { + "text": spliterator.aread_text, + "jsonl": spliterator.aread_jsonl, + "csv": spliterator.aread_csv, + "tsv": spliterator.aread_tsv, + "psv": spliterator.aread_psv, +} + +COUNTERS: dict[str, Callable[..., int]] = { + "text": spliterator.count_text, + "csv": spliterator.count_csv, + "tsv": spliterator.count_tsv, + "psv": spliterator.count_psv, +} + ERRORS = { "UnterminatedQuote": (spliterator.UnterminatedQuoteError, "byteOffset", "byte_offset"), "DuplicateHeader": (spliterator.DuplicateHeaderError, "column", "column"), @@ -62,7 +78,7 @@ def load_cases() -> list[Any]: params.extend( pytest.param(fmt, case, source, id=f"{name}.json :: {case['name']} :: {source}") - for source in ("bytes", "file object", "path") + for source in ("bytes", "file object", "path", "async", "count") ) return params @@ -96,6 +112,16 @@ def test_conformance(fmt: str, case: dict[str, Any], source_kind: str, tmp_path: source = tmp_path / "input" source.write_bytes(data) + if source_kind == "count": + check_count(fmt, case, data, options) + + return + + if source_kind == "async": + check_async(fmt, case, data, options) + + return + rows: list[Any] = [] reader = READERS[fmt](source, **options) @@ -117,3 +143,49 @@ def test_conformance(fmt: str, case: dict[str, Any], source_kind: str, tmp_path: rows = [plain(row) for row in reader] assert rows == case["expect"] + + +def check_count(fmt: str, case: dict[str, Any], data: bytes, options: dict[str, Any]) -> None: + """A count must equal the number of rows the reader yields. It reads no header names, so a repeated header is not + its error to raise.""" + if fmt not in COUNTERS or case.get("error", {}).get("type") == "DuplicateHeader": + pytest.skip("no count for this case") + + if "error" in case: + expected = case["error"] + error_type, json_field, attribute = ERRORS[expected["type"]] + + with pytest.raises(error_type) as raised: + COUNTERS[fmt](data, **options) + + assert getattr(raised.value, attribute) == expected[json_field] + + return + + assert COUNTERS[fmt](data, **options) == len(case["expect"]) + + +def check_async(fmt: str, case: dict[str, Any], data: bytes, options: dict[str, Any]) -> None: + rows: list[Any] = [] + + async def collect() -> None: + async for row in ASYNC_READERS[fmt](data, **options): + rows.append(plain(row)) # noqa: PERF401 - the rows before an error are the point + + if "error" in case: + expected = case["error"] + error_type, json_field, attribute = ERRORS[expected["type"]] + + with pytest.raises(error_type) as raised: + asyncio.run(collect()) + + assert getattr(raised.value, attribute) == expected[json_field] + + if "expectBeforeError" in case: + assert rows == case["expectBeforeError"] + + return + + asyncio.run(collect()) + + assert rows == case["expect"] diff --git a/crates/spliterator-py/tests/test_sources.py b/crates/spliterator-py/tests/test_sources.py index ade7ab1..5b09a17 100644 --- a/crates/spliterator-py/tests/test_sources.py +++ b/crates/spliterator-py/tests/test_sources.py @@ -3,9 +3,11 @@ from __future__ import annotations +import asyncio import gzip import io from pathlib import Path +from typing import Any import pytest import spliterator @@ -103,3 +105,42 @@ def test_bad_options_raise_value_error() -> None: with pytest.raises(ValueError, match="mode"): list(spliterator.read_csv(b"a", mode="rows")) # type: ignore[arg-type] + + +def test_unknown_options_raise_type_error() -> None: + with pytest.raises(TypeError, match="normalise_keys"): + list(spliterator.read_csv(b"a\n1", normalise_keys=True)) # type: ignore[call-arg] + + +def test_async_reader_over_a_compressed_path(tmp_path: Path) -> None: + path = tmp_path / "data.jsonl.gz" + path.write_bytes(gzip.compress(LINES.encode())) + + async def first_three() -> list[Any]: + rows = [] + + async for row in spliterator.aread_jsonl(path): + rows.append(row) + + if len(rows) == 3: + break + + return rows + + assert ids(asyncio.run(first_three())) == [0, 1, 2] + + +def test_async_csv_applies_transformers() -> None: + async def collect() -> list[Any]: + return [row async for row in spliterator.aread_csv(b"n\n1\n2\n", transformers={"n": int})] + + assert asyncio.run(collect()) == [{"n": 1}, {"n": 2}] + + +def test_counts_match_reads(tmp_path: Path) -> None: + path = tmp_path / "part.jsonl.gz" + path.write_bytes(gzip.compress(LINES.encode())) + + assert spliterator.count_text(path) == 5000 + assert spliterator.count_text(path, drop=10, take=7) == 7 + assert spliterator.count_csv(b"k,k\n1,2\n3,4\n") == 2 diff --git a/crates/spliterator/README.md b/crates/spliterator/README.md index ed378a6..4d0d69d 100644 --- a/crates/spliterator/README.md +++ b/crates/spliterator/README.md @@ -24,6 +24,8 @@ assert_eq!(lines, ["a", "b"]); ``` +`text::count` and `csv::count` return the number of rows the matching reader would yield, without building them. + ## Sources `Reader::open` takes a path, bytes, or any `Read` (`Source`). A path `p` reads `p.zst`, or failing that `p.gz`, when one diff --git a/crates/spliterator/src/csv.rs b/crates/spliterator/src/csv.rs index 7f47dbd..1cb8d92 100644 --- a/crates/spliterator/src/csv.rs +++ b/crates/spliterator/src/csv.rs @@ -291,6 +291,50 @@ impl Iterator for Reader { } } +/// The number of data rows [`Reader`] would yield for `source`, without splitting cells. The header is skipped but not +/// read, so a repeated header is not an error here. +pub fn count(source: impl Into, options: Options) -> Result { + count_reader(open(source.into(), options.open)?, options) +} + +/// The number of data rows [`Reader::new`] would yield for an already open, already decompressed stream. +pub fn count_reader(reader: R, options: Options) -> Result { + let records = RecordReader::new( + reader, + RecordOptions { + delimiter: options.row_delimiter.clone(), + enable_quote_handling: options.enable_quote_handling, + allow_unterminated_quote: options.allow_unterminated_quote, + crlf: options.crlf, + skip_empty: options.skip_empty, + }, + )?; + + let mut header_pending = options.header; + let mut window = Window::new(options.drop, options.take); + let mut count = 0; + + for record in records { + if window.exhausted() { + break; + } + + record?; + + if header_pending { + header_pending = false; + + continue; + } + + if window.admit() { + count += 1; + } + } + + Ok(count) +} + /// Splits one record into cells. struct ColumnSplitter { /// The delimiter as text, when it is valid UTF-8; otherwise cells are split on the raw bytes. diff --git a/crates/spliterator/src/text.rs b/crates/spliterator/src/text.rs index 30fe28f..76abde3 100644 --- a/crates/spliterator/src/text.rs +++ b/crates/spliterator/src/text.rs @@ -175,3 +175,49 @@ impl Iterator for Reader { None } } + +/// Whether a raw record would be dropped as blank: decoded, its BOM removed and trimmed to nothing. An ASCII record is +/// judged on its bytes, so only a record holding non-ASCII bytes is decoded. +pub(crate) fn is_blank_record(bytes: &[u8]) -> bool { + for &byte in bytes { + if byte >= 0x80 { + return crate::whitespace::is_blank(&decode_record(bytes)); + } + + if !matches!(byte, b' ' | b'\t' | b'\n' | 0x0B | 0x0C | b'\r') { + return false; + } + } + + true +} + +/// The number of rows [`Reader`] would yield for `source`, without building them. +pub fn count(source: impl Into, options: Options) -> Result { + count_reader(open(source.into(), options.open)?, options) +} + +/// The number of rows [`Reader::new`] would yield for an already open, already decompressed stream. +pub fn count_reader(reader: R, options: Options) -> Result { + let drop_blank = options.trim && options.skip_empty; + let mut window = Window::new(options.drop, options.take); + let mut count = 0; + + for record in RecordReader::new(reader, options.record_options())? { + if window.exhausted() { + break; + } + + let record = record?; + + if drop_blank && is_blank_record(&record.bytes) { + continue; + } + + if window.admit() { + count += 1; + } + } + + Ok(count) +} diff --git a/crates/spliterator/tests/conformance.rs b/crates/spliterator/tests/conformance.rs index c3653c8..6f5541e 100644 --- a/crates/spliterator/tests/conformance.rs +++ b/crates/spliterator/tests/conformance.rs @@ -219,6 +219,19 @@ fn run(format: &str, options: &Map, source: Source) -> (Vec, + source: Source, +) -> Option> { + match format { + "jsonl" => None, + "text" => Some(text::count(source, text_options(options))), + _ => Some(csv::count(source, csv_options(format, options))), + } +} + /// JSON equality with numbers compared by value, so `-2500` from the fixture equals serde's `-2500.0`. fn same(left: &Value, right: &Value) -> bool { match (left, right) { @@ -277,6 +290,24 @@ fn conformance() { let (rows, error) = run(format, &options, source()); let id = format!("{name}.json › {} › {label}", case["name"].as_str().unwrap()); + let expected_error = case.get("error"); + + // A count must equal the rows the reader yields. It reads no header names, so a repeated header is + // not its error to raise. + let duplicate_header = + expected_error.is_some_and(|e| e["type"] == "DuplicateHeader"); + + if let (Some(counted), false) = + (count_rows(format, &options, source()), duplicate_header) + { + checked += 1; + + match (counted, expected_error) { + (Err(error), Some(expected)) if error_matches(&error, expected) => {} + (Ok(n), None) if Some(n) == case["expect"].as_array().map(Vec::len) => {} + (other, _) => failures.push(format!("{id} (count): got {other:?}")), + } + } if let Some(expected) = case.get("error") { match &error { diff --git a/docs/v10.md b/docs/v10.md index 4f81c14..2f34123 100644 --- a/docs/v10.md +++ b/docs/v10.md @@ -203,5 +203,9 @@ for row in spliterator.read_csv("people.csv"): print(row["Full Name"]) ``` -Option names are the TypeScript names, in snake_case for Rust and Python. Neither port offers XLSX, workers, sequences -or `columnScan`; the Rust port has no transformers. +Option names are the TypeScript names, in snake_case for Rust and Python. Text and the CSV family have counts in all +three (`count`/`countAsync`, `text::count`/`csv::count`, `count_text`/`count_csv`/…), and every count equals the number +of rows the matching reader yields; the conformance runners check that on every case. Python's readers each have an async +twin (`aread_text`, `aread_csv`, `aread_jsonl`, …) that reads batches on a worker thread. + +Neither port offers XLSX, workers, sequences or `columnScan`; the Rust port has no transformers. diff --git a/lib/formats/TextSpliterator.ts b/lib/formats/TextSpliterator.ts index a169455..cbbbef6 100644 --- a/lib/formats/TextSpliterator.ts +++ b/lib/formats/TextSpliterator.ts @@ -45,13 +45,17 @@ export interface TextSpliteratorInit { const ASCII_WHITESPACE = new Set([0x20, 0x09, 0x0a, 0x0b, 0x0c, 0x0d]) +const blankDecoder = new TextDecoder() + /** - * Whether a raw row would trim to nothing: what `count` tests so it agrees with `from` without decoding. It checks - * ASCII whitespace, which is what a delimited text file contains; `String.prototype.trim` also strips the rarer Unicode - * spaces, so a row made only of those is counted here and dropped there. + * Whether a raw row would trim to nothing, so `count` agrees exactly with what `from` yields. An ASCII row is judged on + * its bytes; only a row whose first non-whitespace byte is non-ASCII is decoded, since `String.prototype.trim` also + * strips Unicode spaces (U+00A0, U+3000, a BOM). */ function isBlank(row: Uint8Array): boolean { for (const byte of row) { + if (byte >= 0x80) return blankDecoder.decode(row).trim().length === 0 + if (!ASCII_WHITESPACE.has(byte)) return false } @@ -179,9 +183,8 @@ export abstract class TextSpliterator { /** * Count logical text rows without decoding them. * - * This counts the slices {@linkcode from} would yield, including a final unterminated row and its `skipEmpty`, `drop`, - * and `take` behavior. The blank test is ASCII whitespace, so a row of only Unicode spaces is counted here and - * dropped by `from` after decoding. + * This counts the rows {@linkcode from} would yield, including a final unterminated row and its `skipEmpty`, `drop`, + * and `take` behavior. * * @see {@linkcode countAsync} for files and other asynchronous sources. */ @@ -214,9 +217,8 @@ export abstract class TextSpliterator { * Count logical text rows without decoding them. * * This counts the slices {@linkcode fromAsync} would yield, including a final unterminated row and its `skipEmpty`, - * `drop`, and `take` behavior, with the same ASCII-only blank test as {@linkcode count}. A path, URL, or file handle - * is opened independently and can then be passed to {@linkcode fromAsync}; an async chunk source is inherently - * consumed. + * `drop`, and `take` behavior, with the same blank test as {@linkcode count}. A path, URL, or file handle is opened + * independently and can then be passed to {@linkcode fromAsync}; an async chunk source is inherently consumed. */ public static async countAsync( source: AsyncDataResource, diff --git a/test/conformance/conformance.test.ts b/test/conformance/conformance.test.ts index a8f2c44..27e25fb 100644 --- a/test/conformance/conformance.test.ts +++ b/test/conformance/conformance.test.ts @@ -120,6 +120,44 @@ function variants(format: Format, input: ConformanceCase["input"], options: Reco return result } +interface CountVariant { + label: string + run: () => number | Promise +} + +/** + * Every way to count a case's rows. Text and the CSV family have counts; a count must equal the number of rows the + * reader yields, so the expected rows double as the expected count. + */ +function countVariants(format: Format, input: ConformanceCase["input"], options: Record): CountVariant[] { + if (format === "jsonl") return [] + + const counter = READERS[format] as unknown as { + count(source: unknown, init: unknown): number + countAsync(source: unknown, init: unknown): Promise + } + + const bytes = inputBytes(input) + const result: CountVariant[] = [] + + if (typeof input === "string") { + result.push({ label: "count(string)", run: () => counter.count(input, options) }) + } + + result.push({ label: "count(bytes)", run: () => counter.count(bytes, options) }) + + for (const size of [1, Infinity]) { + for (const bulkThreshold of [0, undefined]) { + result.push({ + label: `countAsync, chunks of ${size}, ${bulkThreshold === 0 ? "streamed" : "adaptive"}`, + run: () => counter.countAsync(chunked(bytes, size), { ...options, bulkThreshold }), + }) + } + } + + return result +} + async function collect(run: Variant["run"]): Promise<{ rows: unknown[]; error: unknown }> { const rows: unknown[] = [] @@ -183,6 +221,27 @@ for (const name of FILES) { expect(rows).toEqual(testCase.expect) } ) + + // A count reads no header names, so a repeated header is not its error to raise. + if (testCase.error?.type === "DuplicateHeader") return + + test.for(countVariants(format, testCase.input, testCase.options ?? {}))("$label", async (variant) => { + if (testCase.error) { + let caught: unknown + + try { + await variant.run() + } catch (error) { + caught = error + } + + expectError(caught, testCase.error) + + return + } + + expect(await variant.run()).toBe(testCase.expect!.length) + }) }) } })