import gzip
import operator
from datetime import datetime
from io import StringIO
from typing import IO
from gain import logging
from gain.utils.regions import BedRegion, difference, total_length
from .default_attributes import format_default_attributes
from .gene_models import (
GeneModels,
)
from .transcript_models import (
Exon,
TranscriptModel,
)
logger = logging.getLogger(__name__)
GTF_FEATURE_ORDER: dict[str, int] = {
"gene": 0,
"transcript": 1,
"exon": 2,
"CDS": 2,
"start_codon": 2,
"stop_codon": 2,
"UTR": 3,
}
GTFRecordIndex = tuple[str, int, int, int]
GTFRecord = tuple[GTFRecordIndex, str]
[docs]
def gtf_canonical_index(index: GTFRecordIndex) -> tuple:
# This function converts a GTFRecordIndex for GTF-canonical
# sorting of a GTF file by placing the feature's index at the front
return (index[3], *index[:3])
[docs]
def gene_models_to_gtf(
gene_models: GeneModels, *,
sort_by_position: bool = True,
) -> StringIO:
"""Output a GTF format string representation."""
if not gene_models.gene_models:
logger.warning("Serializing empty (probably not loaded) gene models!")
return StringIO()
record_buffer: list[GTFRecord] = []
for (chrom, gene_name), transcripts in gene_models._chrom_genes(): # ruff: ignore[private-member-access]
t = transcripts[0]
start = min(t.tx[0] for t in transcripts)
stop = max(t.tx[1] for t in transcripts)
strand = t.strand
gene_id = gene_name
version = t.attributes.get("gene_version", ".")
src = t.attributes.get("gene_source", ".")
biotype = t.attributes.get("gene_biotype", ".")
attrs = ";".join([
f'gene_id "{gene_id}"',
f'gene_version "{version}"',
f'gene_name "{gene_name}"',
f'gene_source "{src}"',
f'gene_biotype "{biotype}"',
])
gene_rec = \
f"{chrom}\t{src}\tgene\t{start}\t{stop}\t.\t{strand}\t.\t{attrs};"
record_buffer.append(
((chrom, start, -stop, GTF_FEATURE_ORDER["gene"]), gene_rec))
for transcript in transcripts:
record_buffer.extend(transcript_to_gtf(transcript))
if sort_by_position:
record_buffer.sort(key=operator.itemgetter(0))
else:
record_buffer.sort(key=lambda rec: gtf_canonical_index(rec[0]))
joined_records = "\n".join(rec[1] for rec in record_buffer)
return StringIO(
f"""##description: GTF format dump for gene models "{gene_models.resource.resource_id or '?'}"
##provider: GPF
##format: gtf
##date: {datetime.today().strftime('%Y-%m-%d')}
{joined_records}
""") # ruff: ignore[line-too-long]
[docs]
def get_exon_number_for(
transcript: TranscriptModel,
start: int,
stop: int,
) -> int:
"""Get the exon number for a genomic region.
Returns the exon number (in transcript order) that overlaps the
given genomic coordinates.
Args:
start (int): Start position (1-based).
stop (int): End position (1-based).
Returns:
int: Exon number (1-based) in transcript orientation.
Returns 0 if no overlapping exon found.
Example:
>>> # For a region within the second exon of a + strand transcript
>>> exon_num = transcript.get_exon_number_for(1000, 1050)
>>> print(f"Region is in exon {exon_num}")
Note:
Exon numbering is strand-aware:
- Positive strand: numbered 5' to 3' (exon 1 is first)
- Negative strand: numbered 5' to 3' (exon 1 is last in genome)
"""
for exon_number, exon in enumerate(transcript.exons):
if not (start > exon.stop or stop < exon.start):
return exon_number + 1 if transcript.strand == "+" \
else len(transcript.exons) - exon_number
return 0
[docs]
def build_gtf_record(
transcript: TranscriptModel,
feature: str,
start: int, stop: int,
attrs: str,
) -> GTFRecord:
"""Build an indexed GTF format record for a feature."""
src = transcript.attributes.get("gene_source", ".")
phase = "."
exon_number = -1
if feature in ("exon", "CDS", "start_codon", "stop_codon"):
exon_number = get_exon_number_for(transcript, start, stop)
if feature in ("CDS", "start_codon", "stop_codon"):
frame = calc_frame_for_gtf_cds_feature(
transcript, BedRegion(transcript.chrom, start, stop))
phase = str((3 - frame) % 3)
line = (f"{transcript.chrom}\t{src}\t{feature}\t{start}"
f"\t{stop}\t.\t{transcript.strand}\t{phase}\t{attrs};")
if feature in ("exon", "CDS", "start_codon", "stop_codon"):
line = f'{line}exon_number "{exon_number}";'
# add stop as negative to sort it in descending order
index = \
(transcript.chrom, start, -stop, GTF_FEATURE_ORDER[feature])
return (index, line)
[docs]
def collect_gtf_start_codon_regions(
strand: str,
cds_regions: list[BedRegion],
) -> list[BedRegion]:
"""Returns list of all regions that represent the start codon."""
if strand == "+":
region = cds_regions[0]
if len(region) >= 3:
return [
BedRegion(
region.chrom,
region.start,
region.start + 2,
),
]
result = [region]
for region in cds_regions[1:]:
total = total_length(result)
if total + len(region) >= 3:
result.append(BedRegion(
region.chrom,
region.start,
region.start + (2 - total),
))
return result
result.append(region)
elif strand == "-":
region = cds_regions[-1]
if len(region) >= 3:
return [
BedRegion(
region.chrom,
region.stop - 2,
region.stop,
),
]
result = [region]
for region in reversed(cds_regions[:-1]):
total = total_length(result)
if total + len(region) >= 3:
result.append(BedRegion(
region.chrom,
region.stop - (2 - total),
region.stop,
))
return list(reversed(result))
result.append(region)
else:
raise ValueError("Invalid strand")
return []
[docs]
def collect_gtf_stop_codon_regions(
strand: str,
cds_regions: list[BedRegion],
) -> list[BedRegion]:
"""Returns list of all regions that represent the stop codon."""
if strand == "+":
region = cds_regions[-1]
if len(region) >= 3:
return [
BedRegion(
region.chrom,
region.stop - 2,
region.stop,
),
]
result = [region]
for region in reversed(cds_regions[:-1]):
total = total_length(result)
if total + len(region) >= 3:
result.append(BedRegion(
region.chrom,
region.stop - (2 - total),
region.stop,
))
return list(reversed(result))
result.append(region)
elif strand == "-":
region = cds_regions[0]
if len(region) >= 3:
return [
BedRegion(
region.chrom,
region.start,
region.start + 2,
),
]
result = [region]
for region in cds_regions[1:]:
total = total_length(result)
if total + len(region) >= 3:
result.append(BedRegion(
region.chrom,
region.start,
region.start + (2 - total),
))
return result
result.append(region)
else:
raise ValueError("Invalid strand")
return []
[docs]
def collect_gtf_cds_regions(
strand: str,
cds_regions: list[BedRegion],
) -> list[BedRegion]:
"""Returns list of all regions that represent the CDS."""
stop_codon_regions = collect_gtf_stop_codon_regions(strand, cds_regions)
return difference(cds_regions, stop_codon_regions) # type: ignore
[docs]
def find_exon_cds_region_for_gtf_cds_feature(
transcript: TranscriptModel,
region: BedRegion,
) -> tuple[Exon, BedRegion]:
"""Find exon and CDS region that contains the given feature."""
for exon in transcript.exons:
if exon.contains((region.start, region.stop)):
for cds_region in transcript.cds_regions():
if exon.contains((cds_region.start, cds_region.stop)):
return exon, cds_region
raise ValueError(f"exon for region {region} not found")
[docs]
def calc_frame_for_gtf_cds_feature(
transcript: TranscriptModel,
region: BedRegion,
) -> int:
"""Calculate frame for the given feature."""
exon, cds_region = find_exon_cds_region_for_gtf_cds_feature(
transcript, region)
if exon.frame is None:
raise ValueError(f"frame not found for exon {exon}")
if transcript.strand == "+":
return (exon.frame + (abs(cds_region.start - region.start) % 3)) % 3
return (exon.frame + (abs(cds_region.stop - region.stop) % 3)) % 3
[docs]
def transcript_to_gtf(transcript: TranscriptModel) -> list[GTFRecord]:
"""Output an indexed list of GTF-formatted features of a transcript."""
record_buffer: list[GTFRecord] = []
attributes = {
"transcript_id": transcript.tr_id,
"gene_name": transcript.gene,
"gene_id": transcript.gene,
}
str_attrs = ";".join(f'{k} "{v}"' for k, v in attributes.items())
def write_record(feature: str, start: int, stop: int) -> None:
record_buffer.append(
build_gtf_record(transcript, feature, start, stop, str_attrs))
write_record("transcript", transcript.tx[0], transcript.tx[1])
for exon in transcript.exons:
write_record("exon", exon.start, exon.stop)
if transcript.is_coding():
cds_regions = transcript.cds_regions()
for codon in collect_gtf_start_codon_regions(
transcript.strand, cds_regions):
write_record("start_codon", codon.start, codon.stop)
for cds in collect_gtf_cds_regions(
transcript.strand, cds_regions):
write_record("CDS", cds.start, cds.stop)
for codon in collect_gtf_stop_codon_regions(
transcript.strand, cds_regions):
write_record("stop_codon", codon.start, codon.stop)
for utr in transcript.utr3_regions() + transcript.utr5_regions():
write_record("UTR", utr.start, utr.stop)
return record_buffer
def _no_frame_message(transcript_model: TranscriptModel) -> str:
"""What an exon with no frame is refused with, wherever it is met.
The precondition and the writer's own guard report the same
condition, and which one a caller meets depends only on the entry
point it came through. Spelling the message once is what keeps them
from drifting: they are pinned by different suites, so a reword of
either would otherwise pass unnoticed with both still green.
"""
return (
f"transcript {transcript_model.tr_id} at "
f"{transcript_model.chrom} has an exon with no frame "
f"to write in the exonFrames column; the frames are "
f"computed by update_frames()"
)
def _format_exon_frames(transcript_model: TranscriptModel) -> str:
"""Pack the exon reading frames into the ``exonFrames`` column.
An ``Exon`` built without a frame holds ``None``, which ``str()``
spells as the literal ``None`` -- a token the format cannot express
and the read side refuses, so the record could not be read back
(gain#951). A non-coding exon is ``-1`` here, not ``None``; see
``TranscriptModel.calc_frames``.
Every parser fills the frames in through ``update_frames()`` before
the models escape, so nothing gain has read reaches this. What it
guards is a model built by hand whose frames were never computed,
and it refuses rather than substituting the ``-1`` such an exon
would most likely have taken: writing a value the model never held
is what made the fabricated ``nan`` of gain#931 a defect.
A caller coming through ``save_as_default_gene_models`` no longer
reaches this: the precondition stops it before the output file is
opened (gain#978), because raising from here left a truncated file
behind. This stays as the backstop for the three test modules that
drive the writer directly against an already-open file object. It
guards this one shape only -- a loop over no exons has nothing to
object to, so the exonless transcript the precondition also
refuses passes through here unremarked.
"""
frames = []
for exon in transcript_model.exons:
if exon.frame is None:
raise ValueError(_no_frame_message(transcript_model))
frames.append(str(exon.frame))
return ",".join(frames)
def _save_as_default_gene_models(
gene_models: GeneModels,
outfile: IO,
) -> None:
outfile.write(
"\t".join(
[
"chr",
"trID",
"trOrigId",
"gene",
"strand",
"tsBeg",
"txEnd",
"cdsStart",
"cdsEnd",
"exonStarts",
"exonEnds",
"exonFrames",
"atts",
],
),
)
outfile.write("\n")
for transcript_model in gene_models.transcript_models.values():
exon_starts = ",".join([
str(e.start) for e in transcript_model.exons])
exon_ends = ",".join([
str(e.stop) for e in transcript_model.exons])
exon_frames = _format_exon_frames(transcript_model)
add_atts = format_default_attributes(transcript_model.attributes)
columns = [
transcript_model.chrom,
transcript_model.tr_id,
transcript_model.tr_name,
transcript_model.gene,
transcript_model.strand,
transcript_model.tx[0],
transcript_model.tx[1],
transcript_model.cds[0],
transcript_model.cds[1],
exon_starts,
exon_ends,
exon_frames,
add_atts,
]
# Presence, not truthiness: a coordinate of ``0`` is a value the
# file has to carry, and a truthiness guard wrote it as a blank
# cell that the read side then refused -- gain could not read
# back what gain had written (gain#951).
#
# Nothing reaching here holds ``None`` any more: the exon frame
# was the one optional field in the row, and it is refused
# above. The reads cannot produce one either -- the columnar
# reader passes ``na_filter=False``, so a missing cell arrives
# as the empty string. The branch is kept for the presence
# check itself, which is what this line is about; refusing on
# it as well would be guarding a state the model's own types
# do not offer.
outfile.write(
"\t".join([str(x) if x is not None else "" for x in columns]))
outfile.write("\n")
def _check_default_format_can_express(gene_models: GeneModels) -> None:
"""Refuse models the default format has no way to write down.
Two transcript shapes have no spelling in the columnar format. A
transcript with no exons joins to a blank cell in all three exon
columns, and the read side treats a blank exon cell as a hard parse
error (gain#929) -- so the file gain wrote is one gain cannot read
back, reported through format inference as "can't infer gene models
file format" with the real cause buried in the formats it tried. An
exon with no frame spells `str(None)` into the exonFrames column, a
token the format cannot express either (gain#951).
This runs before the output file is opened, which is the whole
point of it being a separate pass. gain#965 prototyped refusing
from inside the write loop and rejected it: both open branches
create and truncate, so a refusal part-way through leaves a
truncated file holding every record up to the offender -- one that
loads cleanly and is simply missing data. Refusing here creates
nothing and truncates nothing, and a file already at the path is
left as it was.
It does not make the write atomic. A failure during the write
itself -- a full disk, an encoding error -- still leaves a partial
file; covering that would need a write-to-temp-and-rename, which is
a pattern this package does not have.
"""
for transcript_model in gene_models.transcript_models.values():
# Nothing is formatted until there is a message to format: this
# runs over every transcript of every save, and the shape it
# looks for is one no parsed model has.
if not transcript_model.exons:
raise ValueError(
f"transcript {transcript_model.tr_id} at "
f"{transcript_model.chrom} has no exons to write in the "
f"exonStarts, exonEnds and exonFrames columns; the "
f"default format has no spelling for a transcript "
f"without them",
)
if any(exon.frame is None for exon in transcript_model.exons):
raise ValueError(_no_frame_message(transcript_model))
[docs]
def save_as_default_gene_models(
gene_models: GeneModels,
output_filename: str, *,
gzipped: bool = True,
) -> None:
"""Save gene models in a file in default file format."""
_check_default_format_can_express(gene_models)
if gzipped:
if not output_filename.endswith(".gz"):
output_filename = f"{output_filename}.gz"
with gzip.open(output_filename, "wt") as outfile:
_save_as_default_gene_models(gene_models, outfile)
else:
with open(output_filename, "wt") as outfile:
_save_as_default_gene_models(gene_models, outfile)