From 9268c0c5dbf412fa90d857d7be7ffe38c8f0681a Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Tue, 29 Sep 2026 09:50:43 -0400 Subject: [PATCH] Add parsing progress callbacks and preserve quoted attributes --- README.md | 49 +++++++++- gtfparse/__init__.py | 2 +- gtfparse/attribute_parsing.py | 63 ++++++++----- gtfparse/read_gtf.py | 102 ++++++++++++++++----- gtfparse/write_gtf.py | 11 +-- tasks/lessons.md | 6 ++ tasks/todo.md | 80 ++++++++++++++++ tests/test_attribute_values.py | 80 ++++++++++++++++ tests/test_progress.py | 162 +++++++++++++++++++++++++++++++++ tests/test_write_gtf.py | 12 +-- 10 files changed, 503 insertions(+), 64 deletions(-) create mode 100644 tasks/lessons.md create mode 100644 tests/test_attribute_values.py create mode 100644 tests/test_progress.py diff --git a/README.md b/README.md index 8f711ba..d88fd21 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,53 @@ gtfparse ======== Parsing tools for GTF (gene transfer format) files. +## Reporting parsing progress + +Pass `progress_callback(stage, completed, total)` to `read_gtf` to connect +parsing to your own progress UI. The callback runs synchronously; exceptions +from it stop parsing and propagate to the caller. Without a callback, gtfparse +does not display a progress bar or add a progress-library dependency. + +| Stage | Reports | +| --- | --- | +| `read` | `(0, None)` before Polars loads/filters the file; `(rows, rows)` after it succeeds | +| `attributes` | `(0, rows)`, then every 10,000 rows, then `(rows, rows)` | +| `convert` | `(0, None)` before output conversion/transforms; `(rows, rows)` after success | + +Counts refer to rows retained after `features` filtering, not bytes or overall +percentages. `None` means the total is unknown: loading and conversion expose +only stage boundaries, so use an indeterminate indicator during those steps. +Attribute expansion is omitted when `expand_attribute_column=False`. A stage +that fails does not emit completion. An empty filtered result is supported; +its attributes stage emits a single `(0, 0)` event. Empty input files retain +the existing Polars `NoDataError` behavior. + +For example, with the optional `tqdm` package installed: + +```python +from gtfparse import read_gtf +from tqdm.auto import tqdm + +with tqdm(unit="rows") as bar: + def show_progress(stage, completed, total): + if stage != bar.desc: + bar.total = total + bar.reset() + bar.set_description_str(stage) + bar.total = total + bar.update(completed - bar.n) + bar.refresh() + + df = read_gtf("gene_annotations.gtf", progress_callback=show_progress) +``` + +`parse_gtf` reports the `read` stage, `parse_gtf_and_expand_attributes` reports +`read` and `attributes`, `parse_gtf_pandas` reports `read` and `convert`, and +`expand_attribute_strings` reports only `attributes` using the same callback. + +Quoted attribute values preserve spaces, semicolons, and apostrophes. Raw +attributes (`expand_attribute_column=False`) retain their original quote marks. + # Example usage ## Parsing all rows of a GTF file into a Pandas DataFrame @@ -41,5 +88,3 @@ gene_fpkms = { if feature == "gene" } ``` - - diff --git a/gtfparse/__init__.py b/gtfparse/__init__.py index bbf3afb..a3dfadd 100644 --- a/gtfparse/__init__.py +++ b/gtfparse/__init__.py @@ -24,7 +24,7 @@ ) from .write_gtf import write_gtf -__version__ = "2.8.1" +__version__ = "2.9.0" __all__ = [ "GENCODE_BIOTYPE_ALIASES", diff --git a/gtfparse/attribute_parsing.py b/gtfparse/attribute_parsing.py index 9cabfc3..ccea6fb 100644 --- a/gtfparse/attribute_parsing.py +++ b/gtfparse/attribute_parsing.py @@ -11,25 +11,37 @@ # limitations under the License. import logging +import re from collections import OrderedDict from sys import intern logger = logging.getLogger(__name__) +# Quoted values may contain semicolons. Keep this compatible with both Python's +# re module and Polars' regex engine so the low-level split column agrees with +# attribute expansion. Unquoted values and single quotes are accepted as well. +ATTRIBUTE_PATTERN = r"""[^\s;]+[ \t]+(?:"[^"]*"|'[^']*'|[^;]+)""" +_ATTRIBUTE_PAIRS = re.compile(ATTRIBUTE_PATTERN) -def expand_attribute_strings(attribute_strings, quote_char="'", missing_value="", usecols=None): + +def expand_attribute_strings( + attribute_strings, quote_char="'", missing_value="", usecols=None, *, progress_callback=None +): """ The last column of GTF has a variable number of key value pairs of the format: "key1 value1; key2 value2;" Parse these into a dictionary mapping each key onto a list of values, - where the value is None for any row where the key was missing. + using missing_value for any row where the key was missing. Parameters ---------- - attribute_strings : list of str + attribute_strings : sequence of str or sequences of str + Raw attribute fields or already separated key/value pairs per row. + None represents a row with no attributes. quote_char : str - Quote character to remove from values + Additional surrounding quote character to remove from values. + Standard double quotes are always recognized. missing_value : any If an attribute is missing from a row, give it this value. @@ -38,6 +50,11 @@ def expand_attribute_strings(attribute_strings, quote_char="'", missing_value="" If not None, then only expand columns included in this set, otherwise use all columns. + progress_callback : callable, optional + Called as ``callback("attributes", completed_rows, total_rows)`` at + the start, every 10,000 rows, and at completion. An empty input emits + one ``("attributes", 0, 0)`` event. Callback exceptions propagate. + Returns OrderedDict of column->value list mappings, in the order they appeared in the attribute strings. """ @@ -46,29 +63,23 @@ def expand_attribute_strings(attribute_strings, quote_char="'", missing_value="" extra_columns = {} column_order = [] - # - # SOME NOTES ABOUT THE BIZARRE STRING INTERNING GOING ON BELOW - # # While parsing millions of repeated strings (e.g. "gene_id" and "TP53"), # we can save a lot of memory by making sure there's only one string - # object per unique string. The canonical way to do this is using - # the 'intern' function. One problem is that Py2 won't let you intern - # unicode objects, so to get around this we call intern(str(...)). - # - # It also turns out to be faster to check interned strings ourselves - # using a local dictionary, hence the two dictionaries below - # and pair of try/except blocks in the loop. + # object per unique column name. Cache interned names locally as well. column_interned_strings = {} + if progress_callback is not None: + progress_callback("attributes", 0, n) + for i, kv_strings in enumerate(attribute_strings): - if type(kv_strings) is str: - kv_strings = kv_strings.split(";") + if isinstance(kv_strings, str): + kv_strings = _ATTRIBUTE_PAIRS.findall(kv_strings) + elif kv_strings is None: + kv_strings = () for kv in kv_strings: - # We're slicing the first two elements out of split() because - # Ensembl release 79 added values like: - # transcript_support_level "1 (assigned to previous version 5)"; - # ...which gets mangled by splitting on spaces. - parts = kv.strip().split(" ", 2)[:2] + # Split once: values such as transcript_support_level may include + # spaces, and whitespace between a key and its value may vary. + parts = kv.strip().split(None, 1) if len(parts) != 2: continue @@ -78,14 +89,14 @@ def expand_attribute_strings(attribute_strings, quote_char="'", missing_value="" try: column_name = column_interned_strings[column_name] except KeyError: - column_name = intern(str(column_name)) + column_name = intern(column_name) column_interned_strings[column_name] = column_name if usecols is not None and column_name not in usecols: continue - if value[0] == quote_char: - value = value.replace(quote_char, "") + if len(value) >= 2 and value[0] in ('"', quote_char) and value[-1] == value[0]: + value = value[1:-1] try: column = extra_columns[column_name] @@ -102,5 +113,9 @@ def expand_attribute_strings(attribute_strings, quote_char="'", missing_value="" extra_columns[column_name] = column column_order.append(column_name) + completed = i + 1 + if progress_callback is not None and (completed % 10_000 == 0 or completed == n): + progress_callback("attributes", completed, n) + logger.info("Extracted GTF attributes: %s", column_order) return OrderedDict((column_name, extra_columns[column_name]) for column_name in column_order) diff --git a/gtfparse/read_gtf.py b/gtfparse/read_gtf.py index 5633329..6b1f12c 100644 --- a/gtfparse/read_gtf.py +++ b/gtfparse/read_gtf.py @@ -16,7 +16,7 @@ import pandas as pd import polars -from .attribute_parsing import expand_attribute_strings +from .attribute_parsing import ATTRIBUTE_PATTERN, expand_attribute_strings from .parsing_error import ParsingError logger = logging.getLogger(__name__) @@ -108,7 +108,7 @@ def parse_with_polars_lazy( - filepath_or_buffer, split_attributes=True, features=None, fix_quotes_columns=["attribute"] + filepath_or_buffer, split_attributes=True, features=None, fix_quotes_columns=None ): # Categories shares categorical mappings on modern Polars. Older releases # need the global string cache for categoricals from separate reads to match. @@ -120,6 +120,7 @@ def parse_with_polars_lazy( "comment_prefix": "#", "null_values": ".", "schema_overrides": DEFAULT_COLUMN_DTYPES, + "quote_char": None, } try: df = polars.read_csv(filepath_or_buffer, new_columns=REQUIRED_COLUMNS, **kwargs).lazy() @@ -129,11 +130,9 @@ def parse_with_polars_lazy( # Drop empty lines that may appear as all-null rows df = df.filter(polars.col("seqname").is_not_null()) - df = df.with_columns( - [polars.col("frame").fill_null(0), polars.col("attribute").str.replace_all('"', "'")] - ) + df = df.with_columns(polars.col("frame").fill_null(0)) - for fix_quotes_column in fix_quotes_columns: + for fix_quotes_column in fix_quotes_columns or (): # Catch mistaken semicolons by replacing "xyz;" with "xyz" # Required to do this since the Ensembl GTF for Ensembl # release 78 has mistakes such as: @@ -147,28 +146,55 @@ def parse_with_polars_lazy( df = df.filter(polars.col("feature").is_in(features)) if split_attributes: - df = df.with_columns([polars.col("attribute").str.split(";").alias("attribute_split")]) + df = df.with_columns( + polars.col("attribute").str.extract_all(ATTRIBUTE_PATTERN).alias("attribute_split") + ) return df def parse_gtf( - filepath_or_buffer, split_attributes=True, features=None, fix_quotes_columns=["attribute"] + filepath_or_buffer, + split_attributes=True, + features=None, + fix_quotes_columns=None, + *, + progress_callback=None, ): + """Read the fixed GTF columns, preserving raw attribute text. + + If split_attributes is True, add an attribute_split list column with + quoted key/value pairs kept intact. Legacy semicolon cleanup is available + by explicitly setting fix_quotes_columns; it is disabled by default. + progress_callback follows read_gtf's contract and emits only the read stage. + """ + if progress_callback is not None: + progress_callback("read", 0, None) df_lazy = parse_with_polars_lazy( filepath_or_buffer=filepath_or_buffer, split_attributes=split_attributes, features=features, fix_quotes_columns=fix_quotes_columns, ) - return df_lazy.collect() + df = df_lazy.collect() + if progress_callback is not None: + progress_callback("read", len(df), len(df)) + return df def parse_gtf_pandas(*args, **kwargs): - return parse_gtf(*args, **kwargs).to_pandas() + """Like parse_gtf, returning pandas and reporting read/convert progress.""" + df = parse_gtf(*args, **kwargs) + progress_callback = kwargs.get("progress_callback") + if progress_callback is not None: + progress_callback("convert", 0, None) + result = df.to_pandas() + if progress_callback is not None: + progress_callback("convert", len(result), len(result)) + return result def parse_gtf_and_expand_attributes( - filepath_or_buffer, restrict_attribute_columns=None, features=None + filepath_or_buffer, restrict_attribute_columns=None, features=None, *, progress_callback=None ): """ Parse lines into column->values dictionary and then expand @@ -181,26 +207,34 @@ def parse_gtf_and_expand_attributes( ---------- filepath_or_buffer : str or buffer object - chunksize : int - restrict_attribute_columns : list/set of str or None If given, then only use these attribute columns. features : set or None Ignore entries which don't correspond to one of the supplied features + + progress_callback : callable, optional + Follows read_gtf's callback contract, emitting read and attributes stages. """ - df = parse_gtf(filepath_or_buffer=filepath_or_buffer, features=features, split_attributes=True) + df = parse_gtf( + filepath_or_buffer=filepath_or_buffer, + features=features, + split_attributes=False, + progress_callback=progress_callback, + ) if type(restrict_attribute_columns) is str: restrict_attribute_columns = {restrict_attribute_columns} elif restrict_attribute_columns: restrict_attribute_columns = set(restrict_attribute_columns) - df.drop_in_place("attribute") - attribute_pairs = df.drop_in_place("attribute_split") + attributes = df.drop_in_place("attribute") return df.with_columns( [ polars.Series(k, vs) - for (k, vs) in expand_attribute_strings(attribute_pairs).items() - if restrict_attribute_columns is None or k in restrict_attribute_columns + for (k, vs) in expand_attribute_strings( + attributes, + usecols=restrict_attribute_columns, + progress_callback=progress_callback, + ).items() ] ) @@ -275,9 +309,11 @@ def read_gtf( result_type="polars", attribute_aliases=None, cast_version_columns=True, + *, + progress_callback=None, ): """ - Parse a GTF into a dictionary mapping column names to sequences of values. + Parse a GTF into a Polars DataFrame, pandas DataFrame, or dictionary. Parameters ---------- @@ -330,7 +366,20 @@ def read_gtf( `protein_version`, `exon_version`) from strings to pandas nullable Int64 when present. Set to False to keep them as strings. + + progress_callback : callable, optional + Called synchronously as ``callback(stage, completed, total)``. Counts + are rows retained after feature filtering. Stages occur in order: + ``read``, ``attributes`` (if expanded), and ``convert``. Read and convert + emit ``(stage, 0, None)`` before work and ``(stage, rows, rows)`` after + success; None means the total is unknown, so show an indeterminate + indicator. Attribute expansion reports a known total at the start, + every 10,000 rows, and at completion (one event for zero rows). + Exceptions, including those raised by the callback, propagate without + reporting completion of the failed stage. Default: no callback or UI. """ + if result_type not in ("polars", "pandas", "dict"): + raise ValueError("result_type must be one of 'polars', 'pandas', or 'dict'") if type(filepath_or_buffer) is str and not exists(filepath_or_buffer): raise ValueError("GTF file does not exist: %s" % filepath_or_buffer) @@ -352,13 +401,21 @@ def read_gtf( filepath_or_buffer, restrict_attribute_columns=parse_usecols, features=features, + progress_callback=progress_callback, ) else: # When the caller opts out of attribute expansion they want the raw # 'attribute' column verbatim — no need to also produce the # 'attribute_split' helper that parse_gtf adds by default. - result_df = parse_gtf(filepath_or_buffer, features=features, split_attributes=False) + result_df = parse_gtf( + filepath_or_buffer, + features=features, + split_attributes=False, + progress_callback=progress_callback, + ) + if progress_callback is not None: + progress_callback("convert", 0, None) # converting back to pandas here because Polars bugs manifest # as `pyo3_runtime.PanicException: assertion `left == right` failed: impl error` # and are generally insane to chase down @@ -417,11 +474,12 @@ def wrapped_fn(x): valid_columns = [c for c in usecols if c in column_names] result_df = result_df[valid_columns] - result = None if result_type == "pandas": result = result_df elif result_type == "polars": result = polars.from_pandas(result_df) - elif result_type == "dict": + else: result = result_df.to_dict() + if progress_callback is not None: + progress_callback("convert", len(result_df), len(result_df)) return result diff --git a/gtfparse/write_gtf.py b/gtfparse/write_gtf.py index 588a5a6..a404740 100644 --- a/gtfparse/write_gtf.py +++ b/gtfparse/write_gtf.py @@ -142,13 +142,10 @@ def write_gtf( Notes ----- - GTF has no escaping mechanism for the structural characters ``"`` and - ``;`` inside attribute values (nor for the tab and newline that delimit - columns and rows), and read_gtf strips quotes and splits on ``;`` when - parsing. Values returned by read_gtf therefore never contain those - characters, so any DataFrame obtained from read_gtf round-trips exactly. A - DataFrame built by hand whose attribute values contain ``"``, ``;``, a tab, - or a newline cannot be represented losslessly and will not round-trip. + Attribute values are enclosed in double quotes, so spaces, apostrophes, + and semicolons round-trip. This writer does not escape embedded double + quotes, tabs, or newlines; values containing those characters are outside + its round-trip guarantees. Empty and missing attributes are both omitted. """ # Accept a pandas DataFrame too, since read_gtf(result_type="pandas") # returns one; convert to polars so the formatting below is uniform. diff --git a/tasks/lessons.md b/tasks/lessons.md new file mode 100644 index 0000000..364d4bd --- /dev/null +++ b/tasks/lessons.md @@ -0,0 +1,6 @@ +# Lessons + +- When the user narrows the workspace scope, stop sibling-repository work and + keep subsequent fixes, issues, and releases within the named repository. +- Distinguish a shipped warning fix from an unimplemented progress feature; + verify each requested behavior independently before describing release status. diff --git a/tasks/todo.md b/tasks/todo.md index ff980f5..b5b2c81 100644 --- a/tasks/todo.md +++ b/tasks/todo.md @@ -1,5 +1,7 @@ # Issue #77: Polars string-cache compatibility +Released in 2.8.1 via PR #78; the remaining release checklist below is historical. + ## Specification `parse_with_polars_lazy` currently enables Polars' process-wide string cache on @@ -57,3 +59,81 @@ If-Range returned `200`, so the client must restart instead of appending. The probe capped response size to avoid downloading the full genomes. This supports the strong-ETag path in openvax/datacache#80; an ETag is a representation validator, not an independently verified SHA-256 digest. + +# Issue #80: parsing progress and confirmed parsing bugs + +## Specification + +Scope stays within gtfparse. Add a keyword-only `progress_callback(stage, +completed, total)` argument to `read_gtf`, `parse_gtf`, +`parse_gtf_and_expand_attributes`, and `expand_attribute_strings` (and forward +it through `parse_gtf_pandas`). Stages are `read`, `attributes`, and `convert`. +Counts are rows retained after feature filtering. Opaque read/convert work +emits `(stage, 0, None)` before starting and `(stage, rows, rows)` on success; +`None` means an unknown total, not a percentage. Attribute expansion starts +at zero with a known total, emits an update every 10,000 rows, and always +reports completion, including zero rows. No attributes stage when expansion +is disabled. Exceptions (including callback exceptions) propagate immediately; +no false completion on a failed stage. The lazy helper remains lazy and gains +no misleading completion callback. The default remains silent and has no UI +dependency. Document a runnable optional tqdm adapter. + +Fix #23 and #44 together: tokenize quoted key/value pairs without splitting +inside quoted values, retain the complete value after the first whitespace +separator, remove only the surrounding quote pair, and preserve apostrophes, +spaces, and semicolons. Share the token pattern between the direct attribute +helper and Polars' optional `attribute_split` column. Expand raw strings directly +so row updates cover tokenization and expansion together. Preserve missing +attribute and repeated-key semantics. Stop default quote rewriting and +semicolon removal; retain explicitly requested `fix_quotes_columns` cleanup +as an opt-in compatibility option. Raw attributes must retain their source +text. Update writer round-trip tests/docs to reflect supported semicolons. +Follow GENCODE's quoted-key/value format, with unquoted values and single +quotes retained as compatibility extensions; embedded unescaped double quotes +and literal field/line separators remain outside writer round-trip guarantees. + +Reject invalid `result_type` before reading input (#76); leave the unrelated +cleanup in #76 and existing README PR #74 alone. Release as 2.9.0 for the new +optional API. Test empty filtered results and empty attribute lists; empty +files continue to raise the existing Polars error without reporting completion. + +## Plan + +- [x] Read issues, current implementation, tests, and format documentation. +- [x] Create feature branch; check in with the API and parsing plan. +- [x] Implement callback contract and parsing fixes with regression tests. +- [x] Document callback adapter and updated quote/round-trip semantics. +- [x] Compare old/new parsing performance on a representative large fixture; + measure callback overhead and verify results on existing real GTF fixtures. +- [x] Run `./lint.sh` and `./test.sh` with deprecations as errors; review diff. +- [ ] Open PR, pass CI, merge, deploy from clean master, verify PyPI artifacts. + +## Review + +`./lint.sh` passes. `./test.sh -W error::DeprecationWarning` passes all 143 +tests with Polars 1.44.2 and 1.32.0; 1.31.0 passes 142 with the existing +modern-only test skipped. Current coverage is 96%; attribute parsing is 100%. +Regression tests cover progress ordering/counts, bounded updates, all result +types, filtering, raw mode, paths/gzip/text/bytes streams, and cancellation +both at stage boundaries and during expansion. Empty files keep their existing +error and do not emit false completion. The executable README tqdm example +was checked against 995 Ensembl fixture rows. + +The regression run also exposed a null-attribute crash, now tracked as #81. +Missing attribute fields are treated as empty rows during expansion. Existing +expanded outputs from all five real GTF fixtures match 2.8.1 exactly. Tests +demonstrate the intended differences for whitespace, quoted semicolons, +apostrophes, raw attributes, and invalid result types. + +Performance check: repeat the 995 data rows in the Ensembl fixture to 250,000 +rows (69,322,957 bytes); run 2.8.1, the new default, and an event-collecting +callback in fresh Python 3.12 processes with Polars 1.44.2 and four threads. +Three interleaved runs gave wall-time medians of 6.05s, 4.75s, and 4.28s, +respectively, but the full 1.34–6.65s range makes speedup or precise overhead +claims unreliable on this shared machine. All nine runs produced identical +content checksums; callbacks emitted only 30 events for 250,000 rows. A separate +in-process attribute-expansion comparison isolates callback bookkeeping from +file I/O and output conversion; see the PR validation record for its result. + +Release completion is recorded in the PR after CI, merge, and deployment so +the deployment checkout remains clean. diff --git a/tests/test_attribute_values.py b/tests/test_attribute_values.py new file mode 100644 index 0000000..df7e3af --- /dev/null +++ b/tests/test_attribute_values.py @@ -0,0 +1,80 @@ +from io import StringIO + +import pytest + +from gtfparse import expand_attribute_strings, parse_gtf, read_gtf + + +def gtf_line(attributes): + return "1\ttest\tgene\t1\t100\t.\t+\t.\t%s\n" % attributes + + +@pytest.mark.parametrize( + "value", + [ + "two words", + "1 (assigned to previous version 5)", + "part1; part2", + "A;B", + "PRAMEF6;-201", + "O'Brien", + " padded ", + "", + ], +) +def test_quoted_values_survive_all_parsing_paths(value): + attributes = 'description "%s"; gene_id "G1";' % value + expected = {"description": [value], "gene_id": ["G1"]} + assert expand_attribute_strings([attributes]) == expected + parsed = parse_gtf(StringIO(gtf_line(attributes))) + assert parsed["attribute"][0] == attributes + assert expand_attribute_strings(parsed["attribute_split"]) == expected + result = read_gtf(StringIO(gtf_line(attributes))) + assert result["description"].to_list() == [value] + assert result["gene_id"].to_list() == ["G1"] + + +def test_single_quotes_unquoted_values_and_repeated_tags(): + attributes = "gene_id 'G1'; note 'one; two'; count 0; tag basic; tag more words;" + expected = { + "gene_id": ["G1"], + "note": ["one; two"], + "count": ["0"], + "tag": ["basic,more words"], + } + assert expand_attribute_strings([attributes]) == expected + assert read_gtf(StringIO(gtf_line(attributes)))["tag"].to_list() == ["basic,more words"] + + +def test_whitespace_missing_attributes_and_restricted_columns(): + attributes = [None, "", '; malformed; gene_id "G1"; note "x; y";', 'gene_id "G2";'] + assert expand_attribute_strings(attributes, usecols={"gene_id"}) == { + "gene_id": ["", "", "G1", "G2"] + } + assert expand_attribute_strings(['gene_id\t"G1";']) == {"gene_id": ["G1"]} + + +def test_raw_attributes_are_preserved(): + attributes = 'gene_id "G1"; note "O\'Brien; part2;-suffix";' + result = read_gtf(StringIO(gtf_line(attributes)), expand_attribute_column=False) + assert result["attribute"].to_list() == [attributes] + + +def test_missing_attribute_column_values(): + text = gtf_line(".") + gtf_line('gene_id "G1";') + result = read_gtf(StringIO(text)) + assert result["gene_id"].to_list() == ["", "G1"] + + +def test_legacy_semicolon_cleanup_is_explicit(): + attributes = 'gene_name "PRAMEF6;"; transcript_name "PRAMEF6;-201";' + result = parse_gtf(StringIO(gtf_line(attributes)), fix_quotes_columns=["attribute"]) + assert result["attribute"][0] == 'gene_name "PRAMEF6"; transcript_name "PRAMEF6-201";' + + +def test_pre_split_attributes_keep_missing_and_repeated_values(): + rows = [["broken", 'gene_id "G1"', 'tag "first; value"', 'tag "second"'], []] + assert expand_attribute_strings(rows) == { + "gene_id": ["G1", ""], + "tag": ["first; value,second", ""], + } diff --git a/tests/test_progress.py b/tests/test_progress.py new file mode 100644 index 0000000..b47ec75 --- /dev/null +++ b/tests/test_progress.py @@ -0,0 +1,162 @@ +import gzip +from io import BytesIO, StringIO + +import pandas as pd +import polars +import pytest +from polars.testing import assert_frame_equal + +from gtfparse import ( + expand_attribute_strings, + parse_gtf, + parse_gtf_and_expand_attributes, + parse_gtf_pandas, + read_gtf, +) + +GTF = ( + '1\ttest\tgene\t1\t100\t.\t+\t.\tgene_id "G1"; gene_version "2";\n' + '1\ttest\texon\t1\t50\t.\t+\t0\tgene_id "G1"; tag "basic";\n' +) + + +@pytest.mark.parametrize("result_type", ["polars", "pandas", "dict"]) +@pytest.mark.parametrize("expand", [True, False]) +@pytest.mark.parametrize("features,rows", [(None, 2), ({"gene"}, 1), ({"CDS"}, 0)]) +def test_read_progress_preserves_results(result_type, expand, features, rows, capsys): + events = [] + kwargs = { + "result_type": result_type, + "expand_attribute_column": expand, + "features": features, + "attribute_aliases": {"gene_id": "gene"}, + "usecols": ["feature", "start", "gene", "gene_version", "attribute"], + } + expected = read_gtf(StringIO(GTF), **kwargs) + actual = read_gtf( + StringIO(GTF), progress_callback=lambda *event: events.append(event), **kwargs + ) + if result_type == "polars": + assert_frame_equal(actual, expected) + elif result_type == "pandas": + pd.testing.assert_frame_equal(actual, expected) + else: + assert actual == expected + expected_events = [("read", 0, None), ("read", rows, rows)] + if expand: + expected_events += [("attributes", 0, rows)] + if rows: + expected_events.append(("attributes", rows, rows)) + expected_events += [("convert", 0, None), ("convert", rows, rows)] + assert events == expected_events + assert capsys.readouterr() == ("", "") + + +@pytest.mark.parametrize("kind", ["path", "gzip", "bytes", "text"]) +def test_progress_with_supported_inputs(kind, tmp_path): + if kind == "path": + source = tmp_path / "test.gtf" + source.write_text(GTF) + elif kind == "gzip": + source = tmp_path / "test.gtf.gz" + source.write_bytes(gzip.compress(GTF.encode())) + elif kind == "bytes": + source = BytesIO(GTF.encode()) + else: + source = StringIO(GTF) + events = [] + result = read_gtf(source, progress_callback=lambda *event: events.append(event)) + assert result.height == 2 + assert events[-1] == ("convert", 2, 2) + + +@pytest.mark.parametrize( + "parser,stages", + [ + (parse_gtf, ["read"]), + (parse_gtf_and_expand_attributes, ["read", "attributes"]), + (parse_gtf_pandas, ["read", "convert"]), + ], +) +def test_lower_level_progress(parser, stages): + events = [] + parser(StringIO(GTF), progress_callback=lambda *event: events.append(event)) + assert [event[0] for event in events[::2]] == stages + assert all(event[1:] == (2, 2) for event in events[1::2]) + + +@pytest.mark.parametrize("rows", [0, 1, 10_000, 20_003]) +def test_attribute_progress_counts_rows_and_is_bounded(rows): + events = [] + attributes = ['gene_id "G1";'] * rows + result = expand_attribute_strings( + attributes, progress_callback=lambda *event: events.append(event) + ) + assert events[0] == ("attributes", 0, rows) + assert events[-1] == ("attributes", rows, rows) + counts = [event[1] for event in events] + assert counts == sorted(set(counts)) + assert len(events) <= rows // 10_000 + 2 + assert result.get("gene_id", []) == ["G1"] * rows + + +@pytest.mark.parametrize("stage", ["read", "attributes", "convert"]) +def test_callback_errors_propagate_without_completion(stage): + events = [] + error = RuntimeError("cancel parsing") + + def callback(*event): + events.append(event) + if event[0] == stage: + raise error + + source = StringIO(GTF) + with pytest.raises(RuntimeError) as caught: + read_gtf(source, progress_callback=callback) + assert caught.value is error + assert events[-1][0:2] == (stage, 0) + assert not any(event[0] == stage and event[1] == 2 for event in events) + if stage == "read": + assert source.tell() == 0 + + +def test_conversion_errors_do_not_report_completion(): + events = [] + with pytest.raises(ValueError): + read_gtf( + StringIO(GTF), + column_converters={"gene_id": int}, + progress_callback=lambda *event: events.append(event), + ) + assert events[-1] == ("convert", 0, None) + + +def test_callback_can_cancel_during_attribute_expansion(): + events = [] + + def callback(*event): + events.append(event) + if event[:2] == ("attributes", 10_000): + raise RuntimeError("cancel during expansion") + + with pytest.raises(RuntimeError, match="cancel during expansion"): + read_gtf(StringIO(GTF * 10_000), progress_callback=callback) + assert events[-1] == ("attributes", 10_000, 20_000) + assert not any(event[0] == "convert" for event in events) + + +def test_empty_file_keeps_existing_error_without_completion(): + events = [] + with pytest.raises(polars.exceptions.NoDataError): + read_gtf(StringIO(""), progress_callback=lambda *event: events.append(event)) + assert events == [("read", 0, None)] + + +@pytest.mark.parametrize("invalid", ["numpy", "Pandas", "", None, []]) +def test_invalid_result_type_rejected_before_reading(invalid): + source = StringIO(GTF) + events = [] + with pytest.raises(ValueError, match="result_type"): + read_gtf(source, result_type=invalid, progress_callback=lambda *event: events.append(event)) + assert source.tell() == 0 + assert events == [] diff --git a/tests/test_write_gtf.py b/tests/test_write_gtf.py index b1d7c02..9b513b8 100644 --- a/tests/test_write_gtf.py +++ b/tests/test_write_gtf.py @@ -192,17 +192,13 @@ def test_missing_value_fixed_columns_use_dot(tmp_path): assert fields[7] == "." # frame -def test_structural_characters_are_not_round_trippable(tmp_path): - """GTF has no escaping for '"' or ';'; read_gtf strips quotes and splits on - ';'. Pin that such values cannot round-trip so a future change is noticed.""" - df = _minimal_df(gene_id=["A;B"], note=['say "hi"']) +def test_quoted_attribute_values_round_trip(tmp_path): + df = _minimal_df(gene_id=["A;B"], note=["O'Brien; says hello"]) out_path = tmp_path / "special.gtf" write_gtf(df, out_path) recovered = read_gtf(str(out_path)) - # the semicolon split the value apart - assert recovered["gene_id"].to_list() != ["A;B"] - # the double quotes were stripped from the value - assert '"' not in recovered["note"][0] + assert recovered["gene_id"].to_list() == ["A;B"] + assert recovered["note"].to_list() == ["O'Brien; says hello"] def test_header_lines_are_written(tmp_path):