Skip to content

feat(annotator): let subclasses abort VCF annotation #682

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
NopeCharon wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(annotator): let subclasses abort VCF annotation #682

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 82 additions & 59 deletions src/ga4gh/vrs/extras/annotator/vcf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,18 @@ class VcfAnnotatorArgsError(VcfAnnotatorError):
"""Raise for improper args passed to VCF annotator"""


class VcfAnnotatorAbortError(VcfAnnotatorError):
"""Raise to abort annotation of the entire VCF.

By default, an exception raised while processing a VCF record is logged and
recorded in that record's ``VRS_Error`` INFO field, and annotation continues
with the next record. Child classes of ``AbstractVcfAnnotator`` can raise this
exception (e.g. from ``on_vrs_object``) to signal an unrecoverable error, such as
a lost database connection, that should instead stop annotation and propagate
to the caller of ``annotate()``.
"""


class FieldName(str, Enum):
"""Define VCF field names for VRS annotations"""

Expand Down Expand Up @@ -243,6 +255,8 @@ def annotate(
validation checks fail, although all instances of failed validation are
logged as warnings regardless.
:raise VCFAnnotatorError: if no output formats are selected
:raise VcfAnnotatorAbortError: if a child class aborts annotation while a
record is being processed
"""
self.raise_for_output_args(output_vcf_path, **kwargs)
# This can be pushed up to the click arg parsing too
Expand All @@ -264,67 +278,72 @@ def annotate(
if (self.collect_alleles and self.should_collect_alleles(**kwargs))
else None
)
for record in vcf:
if vcf_out:
additional_info_fields = [FieldName.IDS_FIELD]
if vrs_attributes:
additional_info_fields += [
FieldName.STARTS_FIELD,
FieldName.ENDS_FIELD,
FieldName.STATES_FIELD,
FieldName.LENGTHS_FIELD,
FieldName.REPEAT_SUBUNIT_LENGTHS_FIELD,
]
else:
# no INFO field names need to be designated if not producing an annotated VCF
additional_info_fields = []
try:
vrs_field_data = self._get_vrs_data(
record,
allele_collection,
assembly,
additional_info_fields,
vrs_attributes=vrs_attributes,
compute_for_ref=compute_for_ref,
require_validation=require_validation,
**kwargs,
)
except Exception as ex:
_logger.exception("VRS error on %s-%s", record.chrom, record.pos)
err_msg = f"{ex}" or f"{type(ex)}"
err_msg = err_msg.translate(VCF_ESCAPE_MAP)
additional_info_fields = [FieldName.ERROR_FIELD]
vrs_field_data = {FieldName.ERROR_FIELD.value: [err_msg]}

_logger.debug(
"VCF record %s-%s generated vrs_field_data %s",
record.chrom,
record.pos,
vrs_field_data,
)

if output_vcf_path and vcf_out:
for k in additional_info_fields:
# pysam correctly converts None → "." for Integer/Float
# INFO fields, but for String fields it converts None →
# "" (empty bytes), violating the VCF spec. Work around
# by using the literal string "." for String-typed fields.
if k in _STRING_FIELDS:
record.info[k.value] = [
"." if v in ("", None) else v
for v in vrs_field_data[k.value]
]
else:
record.info[k.value] = [
None if v in ("", None) else v
for v in vrs_field_data[k.value]
try:
for record in vcf:
if vcf_out:
additional_info_fields = [FieldName.IDS_FIELD]
if vrs_attributes:
additional_info_fields += [
FieldName.STARTS_FIELD,
FieldName.ENDS_FIELD,
FieldName.STATES_FIELD,
FieldName.LENGTHS_FIELD,
FieldName.REPEAT_SUBUNIT_LENGTHS_FIELD,
]
vcf_out.write(record)

vcf.close()
else:
# no INFO field names need to be designated if not producing an annotated VCF
additional_info_fields = []
try:
vrs_field_data = self._get_vrs_data(
record,
allele_collection,
assembly,
additional_info_fields,
vrs_attributes=vrs_attributes,
compute_for_ref=compute_for_ref,
require_validation=require_validation,
**kwargs,
)
except VcfAnnotatorAbortError:
_logger.exception(
"Aborting VCF annotation at %s-%s", record.chrom, record.pos
)
raise
except Exception as ex:
_logger.exception("VRS error on %s-%s", record.chrom, record.pos)
err_msg = f"{ex}" or f"{type(ex)}"
err_msg = err_msg.translate(VCF_ESCAPE_MAP)
additional_info_fields = [FieldName.ERROR_FIELD]
vrs_field_data = {FieldName.ERROR_FIELD.value: [err_msg]}

_logger.debug(
"VCF record %s-%s generated vrs_field_data %s",
record.chrom,
record.pos,
vrs_field_data,
)

if vcf_out:
vcf_out.close()
if output_vcf_path and vcf_out:
for k in additional_info_fields:
# pysam correctly converts None → "." for Integer/Float
# INFO fields, but for String fields it converts None →
# "" (empty bytes), violating the VCF spec. Work around
# by using the literal string "." for String-typed fields.
if k in _STRING_FIELDS:
record.info[k.value] = [
"." if v in ("", None) else v
for v in vrs_field_data[k.value]
]
else:
record.info[k.value] = [
None if v in ("", None) else v
for v in vrs_field_data[k.value]
]
vcf_out.write(record)
finally:
vcf.close()
if vcf_out:
vcf_out.close()

if allele_collection is not None:
self.on_vrs_object_collection(allele_collection, **kwargs)
Expand All @@ -339,6 +358,10 @@ def on_vrs_object(
Reimplement in a child class to add custom logic. Otherwise, simply pass through
``vrs_allele`` without altering it further or storing it.

Exceptions raised here are recorded in the VCF record's ``VRS_Error`` INFO
field and annotation continues with the next record, unless the exception is
a ``VcfAnnotatorAbortError``, which stops annotation altogether.

:param vcf_coords: CHR-POS-REF-ALT from VCF for this allele
:param vrs_allele: allele translated from coords
:return: final VRS allele
Expand Down
91 changes: 90 additions & 1 deletion tests/extras/test_annotate_vcf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,12 @@
import pytest

from ga4gh.vrs.dataproxy import DataProxyValidationError, SeqRepoRESTDataProxy
from ga4gh.vrs.extras.annotator.vcf import VcfAnnotator, VcfAnnotatorError
from ga4gh.vrs.extras.annotator.vcf import (
VcfAnnotator,
VcfAnnotatorAbortError,
VcfAnnotatorError,
)
from ga4gh.vrs.models import Allele

TEST_DATA_DIR = Path("tests/extras/data")

Expand Down Expand Up @@ -314,3 +319,87 @@ def test_annotate_vcf_rle(vcf_annotator: VcfAnnotator, vcr_cassette, tmp_path: P
assert output_vrs_pkl.exists()
if vcr_cassette.write_protected:
assert vcr_cassette.all_played


class FailingVcfAnnotator(VcfAnnotator):
"""Raise ``error`` when ``on_vrs_object`` reaches the given VCF coordinates"""

def __init__(
self, data_proxy: SeqRepoRESTDataProxy, fail_on: str, error: Exception
) -> None:
"""Initialize the annotator.

:param data_proxy: GA4GH sequence dataproxy instance
:param fail_on: VCF coordinates (CHR-POS-REF-ALT) at which to raise ``error``
:param error: exception to raise
"""
super().__init__(data_proxy)
self.fail_on = fail_on
self.error = error

def on_vrs_object(self, vcf_coords: str, vrs_allele: Allele, **kwargs) -> Allele:
if vcf_coords == self.fail_on:
raise self.error
return super().on_vrs_object(vcf_coords, vrs_allele, **kwargs)


def test_annotate_vcf_record_error(
rest_dataproxy: SeqRepoRESTDataProxy, vcr, tmp_path: Path
):
"""Test that a non-abort error in a child class hook only fails the current
record, and annotation continues with the rest of the VCF
"""
input_vcf = TEST_DATA_DIR / "test_rle.vcf"
output_vcf = tmp_path / "test_record_error_output.vcf"
output_vrs_pkl = tmp_path / "test_record_error_output.pkl"

# test_rle.vcf is the input of test_annotate_vcf_rle, so its cassette covers every
# sequence request made here
with vcr.use_cassette("test_annotate_vcf_rle.yaml"):
vcf_annotator = FailingVcfAnnotator(
rest_dataproxy,
fail_on="1-100210778-AA-AA",
error=RuntimeError("Unable to store allele"),
)
vcf_annotator.annotate(input_vcf, output_vcf, output_pkl_path=output_vrs_pkl)

with pysam.VariantFile(str(output_vcf)) as vcf:
variants = list(vcf)
assert len(variants) == 2
assert variants[0].info["VRS_Error"] == ("Unable to store allele",)
assert "VRS_Allele_IDs" not in variants[0].info
assert "VRS_Error" not in variants[1].info
assert len(variants[1].info["VRS_Allele_IDs"]) == 2
assert output_vrs_pkl.exists()


def test_annotate_vcf_abort(
rest_dataproxy: SeqRepoRESTDataProxy, vcr, tmp_path: Path, caplog
):
"""Test that raising ``VcfAnnotatorAbortError`` from a child class hook stops
annotation of the entire VCF and propagates to the caller
"""
input_vcf = TEST_DATA_DIR / "test_rle.vcf"
output_vcf = tmp_path / "test_abort_output.vcf.gz"
output_vrs_pkl = tmp_path / "test_abort_output.pkl"

with vcr.use_cassette("test_annotate_vcf_rle.yaml"):
vcf_annotator = FailingVcfAnnotator(
rest_dataproxy,
fail_on="1-102995989-CTTT-CTTT",
error=VcfAnnotatorAbortError("Lost connection to storage"),
)
with pytest.raises(VcfAnnotatorAbortError, match="Lost connection to storage"):
vcf_annotator.annotate(
input_vcf, output_vcf, output_pkl_path=output_vrs_pkl
)

assert "Aborting VCF annotation at 1-102995989" in caplog.text
# Collection hooks (e.g. PKL dump) are not run for an aborted annotation
assert not output_vrs_pkl.exists()
# Output VCF is closed, containing only the records annotated before aborting
with pysam.VariantFile(str(output_vcf)) as vcf:
variants = list(vcf)
assert len(variants) == 1
assert variants[0].pos == 100210778
assert len(variants[0].info["VRS_Allele_IDs"]) == 2