fix: tag value types, two char tags, and tests

fgpyo/sam/__init__.py

@@ -846,15 +846,19 @@ def all_recs(self) -> Iterator[AlignedSegment]:
     def set_tag(
         self,
         tag: str,
-        value: Union[str, int, float],
+        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, "Tags must be 2 characters."
+
         for rec in self.all_recs():
             rec.set_tag(tag, value)
 

fgpyo/sam/tests/test_template_iterator.py

@@ -1,3 +1,5 @@
+import pytest
+
 from fgpyo.sam import Template
 from fgpyo.sam.builder import SamBuilder
 
@@ -64,3 +66,31 @@ def test_to_templates() -> None:
     assert len(template2.r2_secondaries) == 0
     assert len(list(template2.primary_recs())) == 1
     assert len(list(template2.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)