diff --git a/README.md b/README.md
index c88d6ac..c5a21ea 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,40 @@ gtf-parser extract genes.gtf -t exon -f gene_id transcript_id exon_number -o exo
gtf-parser extract genes.gtf -t exon -f gene_id --no-dedup
```
+### Build and use an index for faster repeated queries
+
+For large files and repeated lookups (for example repeatedly querying
+`-t gene`), build an index once and reuse it.
+
+```bash
+# Build index (prints output index path)
+gtf-parser index genes.gtf
+
+# Rebuild index after parser/index upgrades or source file changes
+gtf-parser index genes.gtf --force
+
+# Build a smaller/faster index for specific queries
+gtf-parser index genes.gtf -t gene -k gene_id gene_name
+
+# Use index for extraction
+gtf-parser extract genes.gtf -t gene -f gene_id gene_name --index genes.gtf.idx.sqlite
+
+# Use index for BED conversion
+gtf-parser bed genes.gtf gene -i gene_id gene_name --index genes.gtf.idx.sqlite
+
+# Retrieve one specific gene quickly by key
+gtf-parser extract genes.gtf -t gene -f gene_id gene_name --index genes.gtf.idx.sqlite --where gene_id=ENSG00000141510
+
+# Multiple accepted values for the same key
+gtf-parser extract genes.gtf -t gene -f gene_id gene_name --index genes.gtf.idx.sqlite --where gene_id=ENSG1,ENSG2
+
+# Combine filters (AND across keys)
+gtf-parser extract genes.gtf -t transcript -f transcript_id gene_id --index genes.gtf.idx.sqlite --where gene_id=ENSG1 --where transcript_id=ENST1
+```
+
+If an index is stale/incompatible or does not cover requested feature types/keys,
+the CLI prints a warning and falls back to plain file parsing.
+
### Convert to BED
Convert features of a given type to BED6 format. The `name` column is built
@@ -71,12 +105,26 @@ The format is auto-detected from the first data line.
## Python API
```python
-from gtf_parser.parser import parse
+from gtf_parser.parser import build_index, parse
+
+idx = build_index("genes.gtf")
-for record in parse("genes.gtf", feature_types={"gene"}):
+for record in parse("genes.gtf", feature_types={"gene"}, index_path=idx):
print(record.seqname, record.get("gene_id"), record.get("gene_name"))
```
+```python
+from gtf_parser.parser import parse
+
+for record in parse(
+ "genes.gtf",
+ feature_types={"gene"},
+ index_path="genes.gtf.idx.sqlite",
+ attribute_filters={"gene_id": {"ENSG00000141510"}},
+):
+ print(record.get("gene_id"), record.get("gene_name"))
+```
+
```python
from gtf_parser.extract import extract
@@ -93,4 +141,4 @@ to_bed("genes.gtf", feature_type="gene", id_fields=["gene_id", "gene_name"])
- [Edoardo Giacopuzzi](https://github.com/edg1983)
-With help from Claude Opus 4.6.
\ No newline at end of file
+With help from Claude Opus 4.6.
diff --git a/src/gtf_parser/bed.py b/src/gtf_parser/bed.py
index 9d0abf0..3686355 100644
--- a/src/gtf_parser/bed.py
+++ b/src/gtf_parser/bed.py
@@ -2,11 +2,12 @@
from __future__ import annotations
+from collections.abc import Mapping, Set
import sys
from pathlib import Path
from typing import TextIO
-from .parser import Record, parse
+from .parser import parse
def to_bed(
@@ -16,6 +17,8 @@ def to_bed(
id_separator: str = "|",
output: TextIO | None = None,
deduplicate: bool = True,
+ index_path: str | Path | None = None,
+ attribute_filters: Mapping[str, Set[str]] | None = None,
) -> None:
"""Convert GTF/GFF3 records of a given feature type to BED6 format.
@@ -40,11 +43,20 @@ def to_bed(
Writable text stream (default: stdout).
deduplicate:
If ``True``, skip duplicate BED lines.
+ index_path:
+ Optional SQLite index path created with ``gtf-parser index``.
+ attribute_filters:
+ Optional attribute filters in the form ``{key: {value1, value2}}``.
"""
out = output or sys.stdout
seen: set[str] | None = set() if deduplicate else None
- for record in parse(source, feature_types={feature_type}):
+ for record in parse(
+ source,
+ feature_types={feature_type},
+ index_path=index_path,
+ attribute_filters=attribute_filters,
+ ):
name = id_separator.join(record.get(f) or "" for f in id_fields)
# BED is 0-based half-open; GTF/GFF are 1-based inclusive
bed_start = record.start - 1
diff --git a/src/gtf_parser/cli.py b/src/gtf_parser/cli.py
index fc45ea4..b28b188 100644
--- a/src/gtf_parser/cli.py
+++ b/src/gtf_parser/cli.py
@@ -29,18 +29,91 @@ def _add_common_args(sub: argparse.ArgumentParser) -> None:
action="store_true",
help="Disable deduplication of output rows",
)
+ sub.add_argument(
+ "--index",
+ default=None,
+ help="Optional SQLite index path created with `gtf-parser index`",
+ )
+ sub.add_argument(
+ "--where",
+ action="append",
+ default=None,
+ help=(
+ "Attribute filter in the form key=value or key=v1,v2. "
+ "Can be repeated to combine filters."
+ ),
+ )
+
+
+def _parse_where_filters(where_filters: list[str] | None) -> dict[str, set[str]] | None:
+ if not where_filters:
+ return None
+
+ parsed: dict[str, set[str]] = {}
+ for raw_filter in where_filters:
+ key, sep, raw_values = raw_filter.partition("=")
+ if not sep:
+ raise ValueError(f"Invalid --where value '{raw_filter}': expected key=value")
+
+ clean_key = key.strip()
+ if not clean_key:
+ raise ValueError(f"Invalid --where value '{raw_filter}': empty key")
+
+ values = {value.strip() for value in raw_values.split(",") if value.strip()}
+ if not values:
+ raise ValueError(f"Invalid --where value '{raw_filter}': empty value list")
+
+ parsed.setdefault(clean_key, set()).update(values)
+
+ return parsed
+
+
+def _warn_if_index_unusable(
+ source: str,
+ index_path: str | None,
+ feature_types: set[str] | None = None,
+ attribute_filters: dict[str, set[str]] | None = None,
+) -> None:
+ if not index_path:
+ return
+
+ from .parser import index_status
+
+ usable, reason = index_status(
+ source,
+ index_path,
+ feature_types=feature_types,
+ attribute_filters=attribute_filters,
+ )
+ if not usable:
+ print(
+ f"Warning: index '{index_path}' will be ignored ({reason}).",
+ file=sys.stderr,
+ )
def _run_extract(args: argparse.Namespace) -> None:
from .extract import extract
+ try:
+ where_filters = _parse_where_filters(args.where)
+ except ValueError as exc:
+ raise SystemExit(str(exc))
+ feature_types = set(args.feature_types) if args.feature_types else None
+ _warn_if_index_unusable(
+ args.input,
+ args.index,
+ feature_types=feature_types,
+ attribute_filters=where_filters,
+ )
output = open(args.output, "w", encoding="utf-8") if args.output else sys.stdout
try:
- feature_types = set(args.feature_types) if args.feature_types else None
extract(
source=args.input,
fields=args.fields,
feature_types=feature_types,
+ index_path=args.index,
+ attribute_filters=where_filters,
output=output,
separator=args.separator,
deduplicate=not args.no_dedup,
@@ -53,6 +126,16 @@ def _run_extract(args: argparse.Namespace) -> None:
def _run_bed(args: argparse.Namespace) -> None:
from .bed import to_bed
+ try:
+ where_filters = _parse_where_filters(args.where)
+ except ValueError as exc:
+ raise SystemExit(str(exc))
+ _warn_if_index_unusable(
+ args.input,
+ args.index,
+ feature_types={args.feature_type},
+ attribute_filters=where_filters,
+ )
output = open(args.output, "w", encoding="utf-8") if args.output else sys.stdout
try:
to_bed(
@@ -60,6 +143,8 @@ def _run_bed(args: argparse.Namespace) -> None:
feature_type=args.feature_type,
id_fields=args.id_fields,
id_separator=args.id_separator,
+ index_path=args.index,
+ attribute_filters=where_filters,
output=output,
deduplicate=not args.no_dedup,
)
@@ -71,8 +156,13 @@ def _run_bed(args: argparse.Namespace) -> None:
def _run_features(args: argparse.Namespace) -> None:
from .parser import parse
+ try:
+ where_filters = _parse_where_filters(args.where)
+ except ValueError as exc:
+ raise SystemExit(str(exc))
+ _warn_if_index_unusable(args.input, args.index, attribute_filters=where_filters)
features: set[str] = set()
- for record in parse(args.input):
+ for record in parse(args.input, index_path=args.index, attribute_filters=where_filters):
features.add(record.feature)
for f in sorted(features):
print(f)
@@ -81,14 +171,44 @@ def _run_features(args: argparse.Namespace) -> None:
def _run_attributes(args: argparse.Namespace) -> None:
from .parser import parse
+ try:
+ where_filters = _parse_where_filters(args.where)
+ except ValueError as exc:
+ raise SystemExit(str(exc))
feature_types = set(args.feature_types) if args.feature_types else None
+ _warn_if_index_unusable(
+ args.input,
+ args.index,
+ feature_types=feature_types,
+ attribute_filters=where_filters,
+ )
keys: set[str] = set()
- for record in parse(args.input, feature_types=feature_types):
+ for record in parse(
+ args.input,
+ feature_types=feature_types,
+ index_path=args.index,
+ attribute_filters=where_filters,
+ ):
keys.update(record.attributes.keys())
for k in sorted(keys):
print(k)
+def _run_index(args: argparse.Namespace) -> None:
+ from .parser import build_index
+
+ feature_types = set(args.feature_types) if args.feature_types else None
+ attribute_keys = set(args.attribute_keys) if args.attribute_keys else None
+ built = build_index(
+ source=args.input,
+ index_path=args.output,
+ force=args.force,
+ feature_types=feature_types,
+ attribute_keys=attribute_keys,
+ )
+ print(built)
+
+
def main(argv: list[str] | None = None) -> None:
parser = argparse.ArgumentParser(
prog="gtf-parser",
@@ -154,6 +274,20 @@ def main(argv: list[str] | None = None) -> None:
action="store_true",
help="Disable deduplication of output rows",
)
+ p_bed.add_argument(
+ "--index",
+ default=None,
+ help="Optional SQLite index path created with `gtf-parser index`",
+ )
+ p_bed.add_argument(
+ "--where",
+ action="append",
+ default=None,
+ help=(
+ "Attribute filter in the form key=value or key=v1,v2. "
+ "Can be repeated to combine filters."
+ ),
+ )
p_bed.set_defaults(func=_run_bed)
# --- features ---
@@ -162,6 +296,20 @@ def main(argv: list[str] | None = None) -> None:
help="List all distinct feature types present in the file",
)
p_features.add_argument("input", help="Input GTF or GFF3 file")
+ p_features.add_argument(
+ "--index",
+ default=None,
+ help="Optional SQLite index path created with `gtf-parser index`",
+ )
+ p_features.add_argument(
+ "--where",
+ action="append",
+ default=None,
+ help=(
+ "Attribute filter in the form key=value or key=v1,v2. "
+ "Can be repeated to combine filters."
+ ),
+ )
p_features.set_defaults(func=_run_features)
# --- attributes ---
@@ -177,8 +325,55 @@ def main(argv: list[str] | None = None) -> None:
default=None,
help="Restrict to these feature type(s)",
)
+ p_attrs.add_argument(
+ "--index",
+ default=None,
+ help="Optional SQLite index path created with `gtf-parser index`",
+ )
+ p_attrs.add_argument(
+ "--where",
+ action="append",
+ default=None,
+ help=(
+ "Attribute filter in the form key=value or key=v1,v2. "
+ "Can be repeated to combine filters."
+ ),
+ )
p_attrs.set_defaults(func=_run_attributes)
+ # --- index ---
+ p_index = subs.add_parser(
+ "index",
+ help="Build an SQLite index for faster repeated retrieval",
+ )
+ p_index.add_argument("input", help="Input GTF or GFF3 file (plain or .gz)")
+ p_index.add_argument(
+ "-o",
+ "--output",
+ default=None,
+ help="Output index path (default: .idx.sqlite)",
+ )
+ p_index.add_argument(
+ "--force",
+ action="store_true",
+ help="Overwrite existing index file if present",
+ )
+ p_index.add_argument(
+ "-t",
+ "--feature-types",
+ nargs="+",
+ default=None,
+ help="Only index these feature type(s) for faster/smaller indexes",
+ )
+ p_index.add_argument(
+ "-k",
+ "--attribute-keys",
+ nargs="+",
+ default=None,
+ help="Only index these attribute keys for --where lookups",
+ )
+ p_index.set_defaults(func=_run_index)
+
args = parser.parse_args(argv)
args.func(args)
diff --git a/src/gtf_parser/extract.py b/src/gtf_parser/extract.py
index 17ed938..0d01ed6 100644
--- a/src/gtf_parser/extract.py
+++ b/src/gtf_parser/extract.py
@@ -2,12 +2,12 @@
from __future__ import annotations
-import csv
+from collections.abc import Mapping, Set
import sys
from pathlib import Path
-from typing import Iterable, TextIO
+from typing import TextIO
-from .parser import COLUMNS, Record, parse
+from .parser import parse
def extract(
@@ -17,6 +17,8 @@ def extract(
output: TextIO | None = None,
separator: str = "\t",
deduplicate: bool = True,
+ index_path: str | Path | None = None,
+ attribute_filters: Mapping[str, Set[str]] | None = None,
) -> None:
"""Extract selected fields from a GTF/GFF3 file and write them out.
@@ -34,13 +36,22 @@ def extract(
Column separator for the output.
deduplicate:
If ``True``, skip duplicate rows.
+ index_path:
+ Optional SQLite index path created with ``gtf-parser index``.
+ attribute_filters:
+ Optional attribute filters in the form ``{key: {value1, value2}}``.
"""
out = output or sys.stdout
seen: set[tuple[str, ...]] | None = set() if deduplicate else None
out.write(separator.join(fields) + "\n")
- for record in parse(source, feature_types=feature_types):
+ for record in parse(
+ source,
+ feature_types=feature_types,
+ index_path=index_path,
+ attribute_filters=attribute_filters,
+ ):
values = tuple(record.get(f) or "" for f in fields)
if seen is not None:
if values in seen:
diff --git a/src/gtf_parser/parser.py b/src/gtf_parser/parser.py
index 5c46361..f59fd06 100644
--- a/src/gtf_parser/parser.py
+++ b/src/gtf_parser/parser.py
@@ -9,15 +9,21 @@
from __future__ import annotations
-import csv
import gzip
+import json
+import os
import re
+import sqlite3
+from collections.abc import Mapping, Set
from dataclasses import dataclass, field
from pathlib import Path
from typing import IO, Iterator, TextIO
# The 9 standard columns in GTF/GFF files
COLUMNS = ("seqname", "source", "feature", "start", "end", "score", "strand", "frame")
+INDEX_SCHEMA_VERSION = "2"
+_GTF_ATTR_RE = re.compile(r'(\w+)\s+"([^"]*)"')
+_GTF_DETECT_RE = re.compile(r'"\s*;')
@dataclass(slots=True)
@@ -44,7 +50,7 @@ def get(self, key: str) -> str | None:
def _parse_gtf_attributes(raw: str) -> dict[str, str]:
"""Parse GTF-style attributes: key "value"; key2 "value2";"""
attrs: dict[str, str] = {}
- for match in re.finditer(r'(\w+)\s+"([^"]*)"', raw):
+ for match in _GTF_ATTR_RE.finditer(raw):
attrs[match.group(1)] = match.group(2)
return attrs
@@ -62,7 +68,7 @@ def _parse_gff3_attributes(raw: str) -> dict[str, str]:
def _detect_format(attr_string: str) -> str:
"""Detect whether an attribute string is GTF or GFF3 format."""
- if "=" in attr_string and not re.search(r'"\s*;', attr_string):
+ if "=" in attr_string and not _GTF_DETECT_RE.search(attr_string):
return "gff3"
return "gtf"
@@ -75,9 +81,418 @@ def _open_file(path: str | Path) -> IO[str]:
return open(path, encoding="utf-8")
+def default_index_path(source: str | Path) -> Path:
+ """Return the default on-disk index path for an input annotation file."""
+ source = Path(source)
+ return Path(f"{source}.idx.sqlite")
+
+
+def _source_fingerprint(path: Path) -> tuple[str, str]:
+ stat = path.stat()
+ return str(stat.st_size), str(stat.st_mtime_ns)
+
+
+def _parse_meta_set(raw: str | None) -> set[str] | None:
+ """Parse comma-separated metadata sets; ``None`` means unconstrained."""
+ if raw is None or raw == "*":
+ return None
+ return {item for item in raw.split(",") if item}
+
+
+def _index_status(
+ source: Path,
+ index_path: Path,
+ feature_types: set[str] | None = None,
+ attribute_filters: dict[str, set[str]] | None = None,
+) -> tuple[bool, str]:
+ if not index_path.exists():
+ return False, "index file does not exist"
+ if not source.exists():
+ return False, "source file does not exist"
+
+ try:
+ conn = sqlite3.connect(index_path)
+ try:
+ rows = conn.execute("SELECT key, value FROM meta").fetchall()
+ finally:
+ conn.close()
+ except sqlite3.Error as exc:
+ return False, f"cannot read index metadata: {exc}"
+
+ meta = {key: value for key, value in rows}
+ schema = meta.get("schema_version")
+ if schema != INDEX_SCHEMA_VERSION:
+ return False, f"schema mismatch (index={schema}, expected={INDEX_SCHEMA_VERSION})"
+
+ source_size, source_mtime_ns = _source_fingerprint(source)
+ if meta.get("source_path") != str(source.resolve()):
+ return False, "index was built for a different source path"
+ if meta.get("source_size") != source_size or meta.get("source_mtime_ns") != source_mtime_ns:
+ return False, "index is stale (source file changed since index build)"
+
+ indexed_features = _parse_meta_set(meta.get("indexed_feature_types"))
+ if feature_types and indexed_features is not None:
+ missing_features = sorted(feature_types - indexed_features)
+ if missing_features:
+ return (
+ False,
+ "index missing feature type(s): " + ", ".join(missing_features),
+ )
+
+ indexed_attr_keys = _parse_meta_set(meta.get("indexed_attribute_keys"))
+ if attribute_filters and indexed_attr_keys is not None:
+ missing_keys = sorted(set(attribute_filters) - indexed_attr_keys)
+ if missing_keys:
+ return (
+ False,
+ "index missing attribute key(s): " + ", ".join(missing_keys),
+ )
+
+ return True, "ok"
+
+
+def index_status(
+ source: str | Path,
+ index_path: str | Path,
+ feature_types: set[str] | None = None,
+ attribute_filters: Mapping[str, Set[str]] | None = None,
+) -> tuple[bool, str]:
+ """Validate whether *index_path* can be used for *source*."""
+ normalized_filters = _normalize_attribute_filters(attribute_filters)
+ return _index_status(
+ Path(source),
+ Path(index_path),
+ feature_types=feature_types,
+ attribute_filters=normalized_filters,
+ )
+
+
+def _normalize_attribute_filters(
+ attribute_filters: Mapping[str, Set[str]] | None,
+) -> dict[str, set[str]] | None:
+ if not attribute_filters:
+ return None
+
+ normalized: dict[str, set[str]] = {}
+ for key, values in attribute_filters.items():
+ clean_key = key.strip()
+ if not clean_key:
+ continue
+
+ clean_values = {value.strip() for value in values if value.strip()}
+ if not clean_values:
+ continue
+
+ if clean_key in normalized:
+ normalized[clean_key].update(clean_values)
+ else:
+ normalized[clean_key] = clean_values
+
+ return normalized or None
+
+
+def _record_matches_attribute_filters(
+ attributes: dict[str, str],
+ attribute_filters: dict[str, set[str]] | None,
+) -> bool:
+ if not attribute_filters:
+ return True
+ for key, allowed_values in attribute_filters.items():
+ if attributes.get(key) not in allowed_values:
+ return False
+ return True
+
+
+def _parse_stream(
+ fh: TextIO,
+ feature_types: set[str] | None = None,
+ attribute_filters: dict[str, set[str]] | None = None,
+) -> Iterator[Record]:
+ """Parse an already-open GTF/GFF3 text stream."""
+ fmt: str | None = None # auto-detect on first data line
+ parse_attrs = _parse_gtf_attributes
+
+ for line in fh:
+ line = line.rstrip("\n\r")
+ if not line or line.startswith("#"):
+ continue
+
+ parts = line.split("\t", 8)
+ if len(parts) < 9:
+ continue
+
+ feature = parts[2]
+ if feature_types and feature not in feature_types:
+ continue
+
+ attr_raw = parts[8]
+
+ if fmt is None:
+ fmt = _detect_format(attr_raw)
+ parse_attrs = _parse_gff3_attributes if fmt == "gff3" else _parse_gtf_attributes
+
+ attributes = parse_attrs(attr_raw)
+ if not _record_matches_attribute_filters(attributes, attribute_filters):
+ continue
+
+ yield Record(
+ seqname=parts[0],
+ source=parts[1],
+ feature=feature,
+ start=int(parts[3]),
+ end=int(parts[4]),
+ score=parts[5],
+ strand=parts[6],
+ frame=parts[7],
+ attributes=attributes,
+ )
+
+
+def _index_matches_source(
+ source: Path,
+ index_path: Path,
+ feature_types: set[str] | None = None,
+ attribute_filters: dict[str, set[str]] | None = None,
+) -> bool:
+ usable, _ = _index_status(
+ source,
+ index_path,
+ feature_types=feature_types,
+ attribute_filters=attribute_filters,
+ )
+ return usable
+
+
+def _iter_from_index(
+ index_path: Path,
+ feature_types: set[str] | None = None,
+ attribute_filters: dict[str, set[str]] | None = None,
+) -> Iterator[Record]:
+ conn = sqlite3.connect(index_path)
+ try:
+ clauses: list[str] = []
+ params: list[str] = []
+
+ if feature_types:
+ ordered_features = sorted(feature_types)
+ placeholders = ",".join("?" for _ in ordered_features)
+ clauses.append(f"r.feature IN ({placeholders})")
+ params.extend(ordered_features)
+
+ if attribute_filters:
+ subqueries: list[str] = []
+ subquery_params: list[str] = []
+ for key, values in sorted(attribute_filters.items()):
+ ordered_values = sorted(values)
+ placeholders = ",".join("?" for _ in ordered_values)
+ subqueries.append(
+ "SELECT row_num FROM attributes "
+ f"WHERE attr_key = ? AND attr_value IN ({placeholders})"
+ )
+ subquery_params.append(key)
+ subquery_params.extend(ordered_values)
+
+ intersected = " INTERSECT ".join(subqueries)
+ clauses.append(f"r.row_num IN ({intersected})")
+ params.extend(subquery_params)
+
+ where_clause = f" WHERE {' AND '.join(clauses)}" if clauses else ""
+ query = (
+ "SELECT r.seqname, r.source, r.feature, r.start, r.end, r.score, r.strand, r.frame, r.attributes_json "
+ "FROM records r"
+ f"{where_clause} ORDER BY r.row_num"
+ )
+ cursor = conn.execute(query, params)
+
+ for row in cursor:
+ yield Record(
+ seqname=row[0],
+ source=row[1],
+ feature=row[2],
+ start=row[3],
+ end=row[4],
+ score=row[5],
+ strand=row[6],
+ frame=row[7],
+ attributes=json.loads(row[8]),
+ )
+ finally:
+ conn.close()
+
+
+def build_index(
+ source: str | Path,
+ index_path: str | Path | None = None,
+ force: bool = False,
+ feature_types: set[str] | None = None,
+ attribute_keys: set[str] | None = None,
+) -> Path:
+ """Build an SQLite index for fast retrieval.
+
+ Parameters
+ ----------
+ source:
+ Input GTF/GFF3 file path.
+ index_path:
+ Output SQLite path. Defaults to ``.idx.sqlite``.
+ force:
+ Overwrite the output index if it already exists.
+ feature_types:
+ Optional subset of feature types to index.
+ attribute_keys:
+ Optional subset of attribute keys to index for ``attribute_filters``.
+ """
+ source_path = Path(source)
+ if not source_path.exists():
+ raise FileNotFoundError(source_path)
+
+ index = Path(index_path) if index_path is not None else default_index_path(source_path)
+ if index.exists():
+ if force:
+ index.unlink()
+ else:
+ raise FileExistsError(index)
+
+ index.parent.mkdir(parents=True, exist_ok=True)
+ tmp_index = index.with_suffix(index.suffix + ".tmp")
+ if tmp_index.exists():
+ tmp_index.unlink()
+
+ source_size, source_mtime_ns = _source_fingerprint(source_path)
+ normalized_attribute_keys = (
+ {key.strip() for key in attribute_keys if key.strip()} if attribute_keys else None
+ )
+
+ conn: sqlite3.Connection | None = None
+ try:
+ conn = sqlite3.connect(tmp_index)
+ conn.execute("PRAGMA journal_mode=OFF")
+ conn.execute("PRAGMA synchronous=OFF")
+ conn.execute("PRAGMA temp_store=MEMORY")
+ conn.execute(
+ "CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
+ )
+ conn.execute(
+ "CREATE TABLE records ("
+ "row_num INTEGER PRIMARY KEY, "
+ "seqname TEXT NOT NULL, "
+ "source TEXT NOT NULL, "
+ "feature TEXT NOT NULL, "
+ "start INTEGER NOT NULL, "
+ "end INTEGER NOT NULL, "
+ "score TEXT NOT NULL, "
+ "strand TEXT NOT NULL, "
+ "frame TEXT NOT NULL, "
+ "attributes_json TEXT NOT NULL)"
+ )
+ conn.execute(
+ "CREATE TABLE attributes ("
+ "row_num INTEGER NOT NULL, "
+ "attr_key TEXT NOT NULL, "
+ "attr_value TEXT NOT NULL)"
+ )
+ conn.execute("CREATE INDEX idx_records_feature ON records(feature)")
+ conn.execute(
+ "CREATE INDEX idx_attributes_key_value_row "
+ "ON attributes(attr_key, attr_value, row_num)"
+ )
+
+ conn.executemany(
+ "INSERT INTO meta(key, value) VALUES(?, ?)",
+ (
+ ("schema_version", INDEX_SCHEMA_VERSION),
+ ("source_path", str(source_path.resolve())),
+ ("source_size", source_size),
+ ("source_mtime_ns", source_mtime_ns),
+ (
+ "indexed_feature_types",
+ "*" if feature_types is None else ",".join(sorted(feature_types)),
+ ),
+ (
+ "indexed_attribute_keys",
+ "*"
+ if normalized_attribute_keys is None
+ else ",".join(sorted(normalized_attribute_keys)),
+ ),
+ ),
+ )
+
+ row_num = 0
+ batch: list[tuple[int, str, str, str, int, int, str, str, str, str]] = []
+ attr_batch: list[tuple[int, str, str]] = []
+ with _open_file(source_path) as fh:
+ for record in _parse_stream(fh, feature_types=feature_types):
+ row_num += 1
+ batch.append(
+ (
+ row_num,
+ record.seqname,
+ record.source,
+ record.feature,
+ record.start,
+ record.end,
+ record.score,
+ record.strand,
+ record.frame,
+ json.dumps(record.attributes, separators=(",", ":")),
+ )
+ )
+ if normalized_attribute_keys is None:
+ for key, value in record.attributes.items():
+ attr_batch.append((row_num, key, value))
+ else:
+ for key in normalized_attribute_keys:
+ value = record.attributes.get(key)
+ if value is not None:
+ attr_batch.append((row_num, key, value))
+
+ if len(batch) >= 10_000:
+ conn.executemany(
+ "INSERT INTO records("
+ "row_num, seqname, source, feature, start, end, score, strand, frame, attributes_json"
+ ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+ batch,
+ )
+ batch.clear()
+ if len(attr_batch) >= 50_000:
+ conn.executemany(
+ "INSERT INTO attributes(row_num, attr_key, attr_value) VALUES (?, ?, ?)",
+ attr_batch,
+ )
+ attr_batch.clear()
+
+ if batch:
+ conn.executemany(
+ "INSERT INTO records("
+ "row_num, seqname, source, feature, start, end, score, strand, frame, attributes_json"
+ ") VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+ batch,
+ )
+ if attr_batch:
+ conn.executemany(
+ "INSERT INTO attributes(row_num, attr_key, attr_value) VALUES (?, ?, ?)",
+ attr_batch,
+ )
+
+ conn.commit()
+ except Exception:
+ if conn is not None:
+ conn.close()
+ if tmp_index.exists():
+ tmp_index.unlink()
+ raise
+ else:
+ conn.close()
+
+ os.replace(tmp_index, index)
+ return index
+
+
def parse(
source: str | Path | TextIO,
feature_types: set[str] | None = None,
+ index_path: str | Path | None = None,
+ attribute_filters: Mapping[str, Set[str]] | None = None,
) -> Iterator[Record]:
"""Parse a GTF or GFF3 file, yielding Record objects.
@@ -88,49 +503,44 @@ def parse(
feature_types:
If provided, only records whose ``feature`` column matches one of the
given types are yielded. Pass ``None`` to yield all records.
+ index_path:
+ Optional path to an SQLite index created by :func:`build_index`.
+ If provided and valid for *source*, parsing is served from the index.
+ attribute_filters:
+ Optional attribute filters in the form ``{key: {value1, value2}}``.
+ A record matches only if all keys match one of their allowed values.
"""
+ normalized_filters = _normalize_attribute_filters(attribute_filters)
+
if isinstance(source, (str, Path)):
- fh = _open_file(source)
+ source_path = Path(source)
+ if index_path is not None:
+ index = Path(index_path)
+ if _index_matches_source(
+ source_path,
+ index,
+ feature_types=feature_types,
+ attribute_filters=normalized_filters,
+ ):
+ yield from _iter_from_index(
+ index,
+ feature_types=feature_types,
+ attribute_filters=normalized_filters,
+ )
+ return
+
+ fh = _open_file(source_path)
should_close = True
else:
fh = source
should_close = False
- fmt: str | None = None # auto-detect on first data line
-
try:
- for line in fh:
- line = line.rstrip("\n\r")
- if not line or line.startswith("#"):
- continue
-
- parts = line.split("\t")
- if len(parts) < 9:
- continue
-
- feature = parts[2]
- if feature_types and feature not in feature_types:
- continue
-
- attr_raw = parts[8]
-
- if fmt is None:
- fmt = _detect_format(attr_raw)
- parse_attrs = (
- _parse_gff3_attributes if fmt == "gff3" else _parse_gtf_attributes
- )
-
- yield Record(
- seqname=parts[0],
- source=parts[1],
- feature=feature,
- start=int(parts[3]),
- end=int(parts[4]),
- score=parts[5],
- strand=parts[6],
- frame=parts[7],
- attributes=parse_attrs(attr_raw),
- )
+ yield from _parse_stream(
+ fh,
+ feature_types=feature_types,
+ attribute_filters=normalized_filters,
+ )
finally:
if should_close:
fh.close()
diff --git a/tests/test_parser.py b/tests/test_parser.py
index d698dc0..e715b82 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -3,7 +3,9 @@
import io
import textwrap
-from gtf_parser.parser import Record, parse
+import pytest
+
+from gtf_parser.parser import Record, build_index, index_status, parse
from gtf_parser.extract import extract
from gtf_parser.bed import to_bed
@@ -22,6 +24,14 @@
chr1\tensembl\tmRNA\t100\t500\t.\t+\t.\tID=mRNA0001;Parent=gene0001
""")
+GTF_DATA_MULTI_GENE = textwrap.dedent("""\
+ ##description: test
+ chr1\thavana\tgene\t100\t500\t.\t+\t.\tgene_id "G1"; gene_name "ABC";
+ chr1\thavana\ttranscript\t100\t500\t.\t+\t.\tgene_id "G1"; transcript_id "T1"; gene_name "ABC";
+ chr2\thavana\tgene\t10\t90\t.\t-\t.\tgene_id "G2"; gene_name "DEF";
+ chr2\thavana\ttranscript\t10\t90\t.\t-\t.\tgene_id "G2"; transcript_id "T2"; gene_name "DEF";
+""")
+
def test_parse_gtf_all():
records = list(parse(io.StringIO(GTF_DATA)))
@@ -43,6 +53,29 @@ def test_parse_gff3():
assert records[0].get("Name") == "ABC"
+def test_parse_attribute_filters_stream():
+ records = list(
+ parse(
+ io.StringIO(GTF_DATA_MULTI_GENE),
+ feature_types={"gene"},
+ attribute_filters={"gene_id": {"G2"}},
+ )
+ )
+ assert len(records) == 1
+ assert records[0].get("gene_name") == "DEF"
+
+
+def test_parse_attribute_filters_stream_multi_value():
+ records = list(
+ parse(
+ io.StringIO(GTF_DATA_MULTI_GENE),
+ feature_types={"gene"},
+ attribute_filters={"gene_id": {"G1", "G2"}},
+ )
+ )
+ assert len(records) == 2
+
+
def test_extract_dedup():
out = io.StringIO()
extract(io.StringIO(GTF_DATA), fields=["gene_id", "transcript_id"], feature_types={"exon"}, output=out)
@@ -85,3 +118,210 @@ def test_record_get_column():
assert r.get("seqname") == "chr1"
assert r.get("start") == "1"
assert r.get("nonexistent") is None
+
+
+def test_parse_with_index(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ records = list(parse(source, feature_types={"gene"}, index_path=index))
+ assert len(records) == 1
+ assert records[0].feature == "gene"
+ assert records[0].get("gene_id") == "G1"
+
+
+def test_index_status_ok(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ usable, reason = index_status(source, index)
+ assert usable
+ assert reason == "ok"
+
+
+def test_index_status_feature_subset_missing(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source, feature_types={"gene"})
+
+ usable, reason = index_status(source, index, feature_types={"exon"})
+ assert not usable
+ assert "missing feature type" in reason
+
+
+def test_index_status_attribute_subset_missing(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source, attribute_keys={"gene_id"})
+
+ usable, reason = index_status(
+ source,
+ index,
+ attribute_filters={"transcript_id": {"T1"}},
+ )
+ assert not usable
+ assert "missing attribute key" in reason
+
+
+def test_parse_with_index_and_attribute_filters(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA_MULTI_GENE, encoding="utf-8")
+ index = build_index(source)
+
+ records = list(
+ parse(
+ source,
+ feature_types={"gene"},
+ index_path=index,
+ attribute_filters={"gene_id": {"G2"}},
+ )
+ )
+ assert len(records) == 1
+ assert records[0].get("gene_name") == "DEF"
+
+
+def test_parse_fallback_when_feature_not_indexed(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source, feature_types={"gene"})
+
+ records = list(parse(source, feature_types={"exon"}, index_path=index))
+ assert len(records) == 3
+ assert all(record.feature == "exon" for record in records)
+
+
+def test_extract_with_index(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ out = io.StringIO()
+ extract(
+ source,
+ fields=["gene_id", "transcript_id"],
+ feature_types={"exon"},
+ output=out,
+ index_path=index,
+ )
+ lines = out.getvalue().strip().split("\n")
+ assert lines[0] == "gene_id\ttranscript_id"
+ assert len(lines) == 3
+
+
+def test_extract_with_index_and_attribute_filters(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA_MULTI_GENE, encoding="utf-8")
+ index = build_index(source)
+
+ out = io.StringIO()
+ extract(
+ source,
+ fields=["gene_id", "gene_name"],
+ feature_types={"gene"},
+ output=out,
+ index_path=index,
+ attribute_filters={"gene_id": {"G2"}},
+ )
+ lines = out.getvalue().strip().split("\n")
+ assert lines[0] == "gene_id\tgene_name"
+ assert lines[1] == "G2\tDEF"
+ assert len(lines) == 2
+
+
+def test_bed_with_index(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ out = io.StringIO()
+ to_bed(
+ source,
+ feature_type="gene",
+ id_fields=["gene_id", "gene_name"],
+ output=out,
+ index_path=index,
+ )
+ lines = out.getvalue().strip().split("\n")
+ assert len(lines) == 1
+ assert lines[0].split("\t")[3] == "G1|ABC"
+
+
+def test_bed_with_index_and_attribute_filters(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA_MULTI_GENE, encoding="utf-8")
+ index = build_index(source)
+
+ out = io.StringIO()
+ to_bed(
+ source,
+ feature_type="gene",
+ id_fields=["gene_id", "gene_name"],
+ output=out,
+ index_path=index,
+ attribute_filters={"gene_id": {"G2"}},
+ )
+ lines = out.getvalue().strip().split("\n")
+ assert len(lines) == 1
+ assert lines[0].split("\t")[3] == "G2|DEF"
+
+
+def test_index_stale_fallback(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ source.write_text(
+ GTF_DATA + 'chr2\thavana\tgene\t1\t10\t.\t+\t.\tgene_id "G2"; gene_name "DEF";\n',
+ encoding="utf-8",
+ )
+
+ records = list(parse(source, feature_types={"gene"}, index_path=index))
+ assert len(records) == 2
+ assert records[1].get("gene_id") == "G2"
+
+
+def test_index_stale_fallback_with_attribute_filters(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ source.write_text(
+ GTF_DATA + 'chr2\thavana\tgene\t1\t10\t.\t+\t.\tgene_id "G2"; gene_name "DEF";\n',
+ encoding="utf-8",
+ )
+
+ records = list(
+ parse(
+ source,
+ feature_types={"gene"},
+ index_path=index,
+ attribute_filters={"gene_id": {"G2"}},
+ )
+ )
+ assert len(records) == 1
+ assert records[0].get("gene_name") == "DEF"
+
+
+def test_index_status_stale(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+ source.write_text(GTF_DATA_MULTI_GENE, encoding="utf-8")
+
+ usable, reason = index_status(source, index)
+ assert not usable
+ assert "stale" in reason
+
+
+def test_build_index_force(tmp_path):
+ source = tmp_path / "test.gtf"
+ source.write_text(GTF_DATA, encoding="utf-8")
+ index = build_index(source)
+
+ with pytest.raises(FileExistsError):
+ build_index(source, index_path=index)
+
+ rebuilt = build_index(source, index_path=index, force=True)
+ assert rebuilt == index