Add write_gtf: serialize a DataFrame back to GTF (fixes #21)

Adds the inverse of read_gtf. A DataFrame produced by read_gtf -- whether
the attribute column was expanded into per-key columns or left as a raw
'attribute' string -- can be written back out and re-read to recover an
equivalent DataFrame.

This builds on the approach in #46 by Benoitdw, with fixes:

- Python 3.9 compatibility: use typing.Union / typing.Optional instead of
  the `str | Path` union syntax, which requires Python 3.10+ (the project
  targets >=3.9).
- No silent data loss: attribute values are omitted only when None (the
  way read_gtf represents an absent key), not when merely falsy. Values
  like 0 or the empty string are now written and survive a round trip.
- Missing values in the fixed columns are serialized as '.', and null
  scores/frames round-trip correctly.
- The attribute field ends with a trailing ';', matching the canonical
  GTF format emitted by Ensembl/GENCODE and re-parsed by read_gtf.
- A column literally named 'attribute' (unexpanded read) is emitted
  verbatim, so expand_attribute_column=False round-trips too.
- Accepts a pandas DataFrame as well as polars (result_type="pandas").
- Raises ValueError if any required fixed column is missing.

Tests cover round trips for RefSeq and Ensembl fixtures (expanded and
unexpanded), pandas input, header lines, missing-value handling, the
falsy-value regression, and the missing-column error.

Co-authored-by: Benoitdw <bw@oncodna.com>
Claude-Session: https://claude.ai/code/session_014fNxSBm5zvdsDNwN3nhmpS

Author: iskandr

gtfparse/__init__.py

@@ -22,6 +22,7 @@
     parse_gtf_pandas,
     read_gtf,
 )
+from .write_gtf import write_gtf
 
 __version__ = "2.7.1"
 
@@ -37,4 +38,5 @@
     "parse_gtf_and_expand_attributes",
     "parse_gtf_pandas",
     "read_gtf",
+    "write_gtf",
 ]

gtfparse/write_gtf.py

@@ -0,0 +1,138 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from collections.abc import Iterable
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Union
+
+import polars
+
+if TYPE_CHECKING:
+    import pandas
+
+logger = logging.getLogger(__name__)
+
+# The eight tab-separated columns that precede the attribute field of a GTF
+# line, in the order they must be written. Any other column in a DataFrame is
+# treated as an expanded attribute (see read_gtf's expand_attribute_column).
+GTF_FIXED_COLUMNS = [
+    "seqname",
+    "source",
+    "feature",
+    "start",
+    "end",
+    "score",
+    "strand",
+    "frame",
+]
+
+# GTF uses a single dot to denote a missing value in the fixed columns.
+MISSING_VALUE = "."
+
+
+def _format_fixed_value(value) -> str:
+    """Render a fixed-column value, using '.' for missing values."""
+    if value is None:
+        return MISSING_VALUE
+    return str(value)
+
+
+def _format_attributes(row: dict, attribute_columns: list[str], raw_column: Optional[str]) -> str:
+    """
+    Build the GTF attribute field for a single row.
+
+    If ``raw_column`` is given (an unexpanded 'attribute' column) its value is
+    emitted verbatim. Otherwise each expanded attribute column is serialized as
+    ``key "value";`` and the pairs are joined with a single space, matching the
+    format produced by common GTF writers (and parsed by read_gtf).
+
+    Attributes whose value is None are omitted, since that is how read_gtf
+    represents a key that was absent on a given row. Note that a value must be
+    *None* to be skipped -- falsy-but-present values such as 0 or the empty
+    string are written out, so they survive a read/write round trip.
+    """
+    if raw_column is not None:
+        raw = row[raw_column]
+        return "" if raw is None else str(raw)
+
+    parts = []
+    for column_name in attribute_columns:
+        value = row[column_name]
+        if value is None:
+            continue
+        parts.append('%s "%s";' % (column_name, value))
+    return " ".join(parts)
+
+
+def write_gtf(
+    df: Union[polars.DataFrame, "pandas.DataFrame"],
+    path: Union[str, Path],
+    header_lines: Optional[Iterable[str]] = None,
+) -> None:
+    """
+    Write a DataFrame of genomic features back out to a GTF file.
+
+    This is the inverse of :func:`read_gtf`. A DataFrame produced by
+    ``read_gtf`` (whether the attribute column was expanded or not) can be
+    written back out and re-read to recover an equivalent DataFrame.
+
+    Parameters
+    ----------
+    df : polars.DataFrame or pandas.DataFrame
+        Feature rows to write. Must contain the fixed GTF columns
+        (seqname, source, feature, start, end, score, strand, frame).
+        Any additional columns are written as attributes, except a column
+        literally named 'attribute', which is treated as a pre-formatted
+        attribute string and emitted verbatim.
+
+    path : str or pathlib.Path
+        Destination file path. Any existing file is overwritten.
+
+    header_lines : iterable of str, optional
+        Lines to write at the top of the file before any feature rows, e.g.
+        ``["##description: example", "##provider: GENCODE"]``. Each is written
+        verbatim on its own line, so include a leading '#' if you want it to be
+        parsed back as a comment.
+    """
+    # Accept a pandas DataFrame too, since read_gtf(result_type="pandas")
+    # returns one; convert to polars so the row iteration below is uniform.
+    if not isinstance(df, polars.DataFrame):
+        df = polars.from_pandas(df)
+
+    columns = df.columns
+    fixed_columns = [name for name in GTF_FIXED_COLUMNS if name in columns]
+    missing_fixed = [name for name in GTF_FIXED_COLUMNS if name not in columns]
+    if missing_fixed:
+        raise ValueError(
+            "DataFrame is missing required GTF column(s): %s" % ", ".join(missing_fixed)
+        )
+
+    # A column named 'attribute' is the raw, unexpanded attribute string;
+    # everything else that isn't a fixed column is an expanded attribute.
+    raw_column = "attribute" if "attribute" in columns else None
+    attribute_columns = [
+        name for name in columns if name not in GTF_FIXED_COLUMNS and name != "attribute"
+    ]
+
+    n_rows = 0
+    with open(path, "w") as output_file:
+        if header_lines is not None:
+            for line in header_lines:
+                output_file.write("%s\n" % line)
+        for row in df.iter_rows(named=True):
+            fixed_fields = [_format_fixed_value(row[name]) for name in fixed_columns]
+            attribute_field = _format_attributes(row, attribute_columns, raw_column)
+            output_file.write("%s\t%s\n" % ("\t".join(fixed_fields), attribute_field))
+            n_rows += 1
+
+    logger.info("Wrote %d GTF rows to %s", n_rows, path)

tests/test_write_gtf.py

@@ -0,0 +1,112 @@
+import polars
+import pytest
+from polars.testing import assert_frame_equal
+
+from gtfparse import read_gtf, write_gtf
+
+from .data import data_path
+
+REFSEQ_GTF_PATH = data_path("refseq.ucsc.small.gtf")
+ENSEMBL_GTF_PATH = data_path("ensembl_grch37.head.gtf")
+
+
+def _assert_round_trips(source_path, tmp_path, expand_attribute_column=True):
+    """read -> write -> read should recover an equivalent DataFrame."""
+    original = read_gtf(source_path, expand_attribute_column=expand_attribute_column)
+    out_path = tmp_path / "round_trip.gtf"
+    write_gtf(original, out_path)
+    recovered = read_gtf(str(out_path), expand_attribute_column=expand_attribute_column)
+    # categorical_as_str avoids spurious mismatches from category orderings
+    assert_frame_equal(original, recovered, categorical_as_str=True)
+
+
+def test_round_trip_expanded_refseq(tmp_path):
+    _assert_round_trips(REFSEQ_GTF_PATH, tmp_path, expand_attribute_column=True)
+
+
+def test_round_trip_expanded_ensembl(tmp_path):
+    # Ensembl file exercises many attribute columns and null score values
+    _assert_round_trips(ENSEMBL_GTF_PATH, tmp_path, expand_attribute_column=True)
+
+
+def test_round_trip_unexpanded(tmp_path):
+    # A raw, unexpanded 'attribute' column should be written verbatim
+    _assert_round_trips(REFSEQ_GTF_PATH, tmp_path, expand_attribute_column=False)
+
+
+def test_round_trip_from_pandas(tmp_path):
+    # write_gtf should also accept a pandas DataFrame (result_type="pandas")
+    original_polars = read_gtf(ENSEMBL_GTF_PATH)
+    original_pandas = read_gtf(ENSEMBL_GTF_PATH, result_type="pandas")
+    out_path = tmp_path / "from_pandas.gtf"
+    write_gtf(original_pandas, out_path)
+    recovered = read_gtf(str(out_path))
+    assert_frame_equal(original_polars, recovered, categorical_as_str=True)
+
+
+def test_falsy_attribute_values_are_preserved(tmp_path):
+    """Attributes whose value is 0 or "" must be written; only None is omitted."""
+    df = polars.DataFrame(
+        {
+            "seqname": ["chr1"],
+            "source": ["test"],
+            "feature": ["gene"],
+            "start": [1],
+            "end": [100],
+            "score": [None],
+            "strand": ["+"],
+            "frame": [None],
+            "gene_id": ["G1"],
+            "zero_attr": ["0"],
+            "empty_attr": [""],
+            "missing_attr": [None],
+        }
+    )
+    out_path = tmp_path / "falsy.gtf"
+    write_gtf(df, out_path)
+    line = out_path.read_text().strip()
+    assert 'zero_attr "0"' in line
+    assert 'empty_attr ""' in line
+    # a None-valued attribute is omitted entirely
+    assert "missing_attr" not in line
+
+
+def test_missing_value_fixed_columns_use_dot(tmp_path):
+    """None values in the fixed columns are serialized as '.'."""
+    df = polars.DataFrame(
+        {
+            "seqname": ["chr1"],
+            "source": ["test"],
+            "feature": ["gene"],
+            "start": [1],
+            "end": [100],
+            "score": [None],
+            "strand": ["+"],
+            "frame": [None],
+            "gene_id": ["G1"],
+        }
+    )
+    out_path = tmp_path / "dots.gtf"
+    write_gtf(df, out_path)
+    fields = out_path.read_text().strip().split("\t")
+    assert fields[5] == "."  # score
+    assert fields[7] == "."  # frame
+
+
+def test_header_lines_are_written(tmp_path):
+    df = read_gtf(REFSEQ_GTF_PATH)
+    out_path = tmp_path / "with_header.gtf"
+    header = ["##description: test", "##provider: gtfparse"]
+    write_gtf(df, out_path, header_lines=header)
+    lines = out_path.read_text().splitlines()
+    assert lines[0] == "##description: test"
+    assert lines[1] == "##provider: gtfparse"
+    # header comment lines are ignored by read_gtf, so it still round-trips
+    recovered = read_gtf(str(out_path))
+    assert_frame_equal(df, recovered, categorical_as_str=True)
+
+
+def test_missing_required_column_raises(tmp_path):
+    df = polars.DataFrame({"seqname": ["chr1"], "gene_id": ["G1"]})
+    with pytest.raises(ValueError, match="missing required GTF column"):
+        write_gtf(df, tmp_path / "bad.gtf")