Add method purge_contexts to CoverageData

coverage/sqldata.py

@@ -641,6 +641,31 @@ def purge_files(self, filenames: Collection[str]) -> None:
                     continue
                 con.execute_void(sql, (file_id,))
 
+    def purge_contexts(self, contexts: Collection[str]) -> None:
+        """Purge any existing coverage data for the given `contexts`.
+
+        This removes all coverage data for the contexts, but does not remove
+        them from the list returned by measured_contexts(), so that context_ids
+        for the contexts remain constant over time."""
+
+        if self._debug.should("dataop"):
+            self._debug.write("Purging contexts {contexts}")
+        self._start_using()
+        with self._connect() as con:
+
+            if self._has_lines:
+                sql = "delete from line_bits where context_id=?"
+            elif self._has_arcs:
+                sql = "delete from arc where context_id=?"
+            else:
+                raise DataError("Can't purge contexts in an empty CoverageData")
+
+            for context in contexts:
+                context_id = self._context_id(context)
+                if context_id is None:
+                    continue
+                con.execute_void(sql, (context_id,))
+
     def update(self, other_data: CoverageData, aliases: Optional[PathAliases] = None) -> None:
         """Update this data with data from several other :class:`CoverageData` instances.
 

tests/test_data.py

@@ -617,6 +617,38 @@ def test_cant_purge_in_empty_data(self) -> None:
         with pytest.raises(DataError, match=msg):
             covdata.purge_files(["abc.py"])
 
+    def test_purge_contexts_lines(self) -> None:
+        covdata = DebugCoverageData()
+        covdata.set_context('context1')
+        covdata.add_lines(LINES_1)
+        covdata.set_context('context2')
+        covdata.add_lines(LINES_2)
+        assert_line_counts(covdata, SUMMARY_1_2)
+        covdata.purge_contexts(["context2"])
+        assert_line_counts(covdata, {'a.py': 2, 'b.py': 1, 'c.py': 0})
+        covdata.purge_contexts(["context1"])
+        assert_line_counts(covdata, {"a.py": 0, "b.py": 0, "c.py": 0})
+        # It's OK to "purge" a context that wasn't measured.
+        covdata.purge_contexts(["context3"])
+        assert_line_counts(covdata, {"a.py": 0, "b.py": 0, "c.py": 0})
+
+    def test_purge_contexts_arcs(self) -> None:
+        covdata = CoverageData()
+        covdata.set_context('context1')
+        covdata.add_arcs(ARCS_3)
+        covdata.set_context('context2')
+        covdata.add_arcs(ARCS_4)
+        assert_line_counts(covdata, SUMMARY_3_4)
+        covdata.purge_contexts(["context1"])
+        assert_line_counts(covdata,{'x.py': 2, 'y.py': 0, 'z.py': 1})
+        covdata.purge_contexts(["context2"])
+        assert_line_counts(covdata, {"x.py": 0, "y.py": 0, "z.py": 0})
+
+    def test_cant_purge_contexts_in_empty_data(self) -> None:
+        covdata = DebugCoverageData()
+        msg = "Can't purge contexts in an empty CoverageData"
+        with pytest.raises(DataError, match=msg):
+            covdata.purge_contexts(["context1"])
 
 class CoverageDataInTempDirTest(CoverageTest):
     """Tests of CoverageData that need a temporary directory to make files."""