Cigar tools: add getters for coordinate positions

cfe-lab · Nov 8, 2023 · eeec6a5 · eeec6a5
1 parent 9073573
commit eeec6a5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/micall/tests/test_cigar_tools.py b/micall/tests/test_cigar_tools.py
@@ -62,7 +62,7 @@ def test_cigar_to_coordinate_mapping(cigar_str, expected_mapping):
     mapping = Cigar.coerce(cigar_str).coordinate_mapping
 
     assert expected_mapping == mapping.ref_to_query_d
-    assert expected_mapping == {i: mapping.ref_to_query(i) for i in mapping.ref_to_query_d}
+    assert expected_mapping == {i: mapping.ref_to_query(i) for i in mapping.reference_coordinates()}
 
 
 @pytest.mark.parametrize("cigar_str", [x[0] for x in cigar_mapping_cases])
@@ -87,7 +87,8 @@ def test_cigar_to_closest_coordinate_mapping(cigar_str, expected_closest_mapping
 
     else:
         fullrange = {i: mapping.ref_to_closest_query(i) \
-                     for i in range(min(mapping.ref_to_query_d), 1 + max(mapping.ref_to_query_d))}
+                     for i in range(min(mapping.reference_coordinates()),
+                                    1 + max(mapping.reference_coordinates()))}
         assert expected_closest_mapping == fullrange
 
 
@@ -102,7 +103,7 @@ def test_cigar_hit_to_coordinate_mapping(cigar_str, expected_mapping):
     assert mapping.ref_to_query(0) == None
     assert mapping.query_to_ref(0) == None
     assert expected_mapping \
-        == {i: mapping.ref_to_query(i) for i in mapping.ref_to_query_d}
+        == {i: mapping.ref_to_query(i) for i in mapping.reference_coordinates()}
 
 
 @pytest.mark.parametrize("cigar_str, expected_closest_mapping", [(x[0], x[2]) for x in cigar_mapping_cases])
@@ -119,7 +120,8 @@ def test_cigar_hit_to_coordinate_closest_mapping(cigar_str, expected_closest_map
         # Coordinates are translated by q_st and r_st.
         expected_closest_mapping = {k + hit.r_st: v + hit.q_st for (k, v) in expected_closest_mapping.items()}
         fullrange = {i: mapping.ref_to_closest_query(i) \
-                     for i in range(min(mapping.ref_to_query_d), 1 + max(mapping.ref_to_query_d))}
+                     for i in range(min(mapping.reference_coordinates()),
+                                    1 + max(mapping.reference_coordinates()))}
         assert expected_closest_mapping == fullrange
 
 

diff --git a/micall/utils/cigar_tools.py b/micall/utils/cigar_tools.py
@@ -60,6 +60,14 @@ def extend(self,
                 self.query_to_op_d[query_index] = op_index
 
 
+    def reference_coordinates(self):
+        return self.ref_to_query_d.keys()
+
+
+    def query_coordinates(self):
+        return self.query_to_ref_d.keys()
+
+
     def ref_to_query(self, index) -> Optional[int]:
         return self.ref_to_query_d.get(index, None)