diff --git a/fgpyo/sam/__init__.py b/fgpyo/sam/__init__.py index 49e27f2d..e054350a 100644 --- a/fgpyo/sam/__init__.py +++ b/fgpyo/sam/__init__.py @@ -863,6 +863,25 @@ def write_to( for rec in rec_iter: writer.write(rec) + def set_tag( + self, + tag: str, + value: Union[str, int, float, None], + ) -> None: + """Add a tag to all records associated with the template. + + Setting a tag to `None` will remove the tag. + + Args: + tag: The name of the tag. + value: The value of the tag. + """ + + assert len(tag) == 2, f"Tags must be 2 characters: {tag}." + + for rec in self.all_recs(): + rec.set_tag(tag, value) + class TemplateIterator(Iterator[Template]): """ diff --git a/fgpyo/sam/tests/test_template_iterator.py b/fgpyo/sam/tests/test_template_iterator.py index cae2f328..79ab9283 100644 --- a/fgpyo/sam/tests/test_template_iterator.py +++ b/fgpyo/sam/tests/test_template_iterator.py @@ -1,5 +1,7 @@ from pathlib import Path +import pytest + from fgpyo.sam import Template from fgpyo.sam import TemplateIterator from fgpyo.sam import reader @@ -99,3 +101,31 @@ def test_write_template( with reader(bam_path) as bam_reader: template = next(TemplateIterator(bam_reader)) assert len([r for r in template.all_recs()]) == 2 + + +def test_set_tag() -> None: + builder = SamBuilder() + template = Template.build(builder.add_pair(chrom="chr1", start1=100, start2=200)) + + TAG = "XF" + VALUE = "value" + + for read in template.all_recs(): + with pytest.raises(KeyError): + read.get_tag(TAG) + + # test setting + template.set_tag(TAG, VALUE) + assert template.r1.get_tag(TAG) == VALUE + assert template.r2.get_tag(TAG) == VALUE + + # test removal + template.set_tag(TAG, None) + for read in template.all_recs(): + with pytest.raises(KeyError): + read.get_tag(TAG) + + # test tags that aren't two characters + for bad_tag in ["", "A", "ABC", "ABCD"]: + with pytest.raises(AssertionError, match="Tags must be 2 characters"): + template.set_tag(bad_tag, VALUE)