Combining contexts works

Author: nedbat

coverage/sqldata.py

@@ -12,6 +12,7 @@
 # TODO: run_info
 
 import glob
+import itertools
 import os
 import sqlite3
 
@@ -90,7 +91,7 @@ def __init__(self, basename=None, suffix=None, warn=None, debug=None):
         self._has_lines = False
         self._has_arcs = False
 
-        self._context_id = 0
+        self._current_context_id = None
 
     def _choose_filename(self):
         self.filename = self._basename
@@ -104,6 +105,7 @@ def _reset(self):
         self._db = None
         self._file_map = {}
         self._have_used = False
+        self._current_context_id = None
 
     def _create_db(self):
         if self._debug and self._debug.should('dataio'):
@@ -178,17 +180,28 @@ def _file_id(self, filename, add=False):
                     self._file_map[filename] = cur.lastrowid
         return self._file_map.get(filename)
 
+    def _context_id(self, context):
+        """Get the id for a context."""
+        assert context is not None
+        self._start_using()
+        with self._connect() as con:
+            row = con.execute("select id from context where context = ?", (context,)).fetchone()
+            if row is not None:
+                return row[0]
+            else:
+                return None
+
     def set_context(self, context):
         """Set the current context for future `add_lines` etc."""
         self._start_using()
         context = context or ""
         with self._connect() as con:
             row = con.execute("select id from context where context = ?", (context,)).fetchone()
             if row is not None:
-                self._context_id = row[0]
+                self._current_context_id = row[0]
             else:
                 cur = con.execute("insert into context (context) values (?)", (context,))
-                self._context_id = cur.lastrowid
+                self._current_context_id = cur.lastrowid
 
     def add_lines(self, line_data):
         """Add measured line data.
@@ -204,10 +217,12 @@ def add_lines(self, line_data):
             ))
         self._start_using()
         self._choose_lines_or_arcs(lines=True)
+        if self._current_context_id is None:
+            self.set_context("")
         with self._connect() as con:
             for filename, linenos in iitems(line_data):
                 file_id = self._file_id(filename, add=True)
-                data = [(file_id, self._context_id, lineno) for lineno in linenos]
+                data = [(file_id, self._current_context_id, lineno) for lineno in linenos]
                 con.executemany(
                     "insert or ignore into line (file_id, context_id, lineno) values (?, ?, ?)",
                     data,
@@ -227,10 +242,12 @@ def add_arcs(self, arc_data):
             ))
         self._start_using()
         self._choose_lines_or_arcs(arcs=True)
+        if self._current_context_id is None:
+            self.set_context("")
         with self._connect() as con:
             for filename, arcs in iitems(arc_data):
                 file_id = self._file_id(filename, add=True)
-                data = [(file_id, self._context_id, fromno, tono) for fromno, tono in arcs]
+                data = [(file_id, self._current_context_id, fromno, tono) for fromno, tono in arcs]
                 con.executemany(
                     "insert or ignore into arc (file_id, context_id, fromno, tono) values (?, ?, ?, ?)",
                     data,
@@ -306,19 +323,23 @@ def update(self, other_data, aliases=None):
 
         # lines
         if other_data._has_lines:
-            for filename in other_data.measured_files():
-                lines = set(other_data.lines(filename))
-                filename = aliases.map(filename)
-                lines.update(self.lines(filename) or ())
-                self.add_lines({filename: lines})
+            for context in other_data.measured_contexts():
+                self.set_context(context)
+                for filename in other_data.measured_files():
+                    lines = set(other_data.lines(filename, context=context))
+                    filename = aliases.map(filename)
+                    lines.update(self.lines(filename, context=context) or ())
+                    self.add_lines({filename: lines})
 
         # arcs
         if other_data._has_arcs:
-            for filename in other_data.measured_files():
-                arcs = set(other_data.arcs(filename))
-                filename = aliases.map(filename)
-                arcs.update(self.arcs(filename) or ())
-                self.add_arcs({filename: arcs})
+            for context in other_data.measured_contexts():
+                self.set_context(context)
+                for filename in other_data.measured_files():
+                    arcs = set(other_data.arcs(filename, context=context))
+                    filename = aliases.map(filename)
+                    arcs.update(self.arcs(filename, context=context) or ())
+                    self.add_arcs({filename: arcs})
 
         # file_tracers
         for filename in other_data.measured_files():
@@ -407,12 +428,11 @@ def file_tracer(self, filename):
                 return row[0] or ""
             return ""   # File was measured, but no tracer associated.
 
-    def lines(self, filename):
+    def lines(self, filename, context=None):
         self._start_using()
         if self.has_arcs():
-            arcs = self.arcs(filename)
+            arcs = self.arcs(filename, context=context)
             if arcs is not None:
-                import itertools
                 all_lines = itertools.chain.from_iterable(arcs)
                 return list(set(l for l in all_lines if l > 0))
 
@@ -421,18 +441,28 @@ def lines(self, filename):
             if file_id is None:
                 return None
             else:
-                linenos = con.execute("select lineno from line where file_id = ?", (file_id,))
+                query = "select lineno from line where file_id = ?"
+                data = [file_id]
+                if context is not None:
+                    query += " and context_id = ?"
+                    data += [self._context_id(context)]
+                linenos = con.execute(query, data)
                 return [lineno for lineno, in linenos]
 
-    def arcs(self, filename):
+    def arcs(self, filename, context=None):
         self._start_using()
         with self._connect() as con:
             file_id = self._file_id(filename)
             if file_id is None:
                 return None
             else:
-                arcs = con.execute("select fromno, tono from arc where file_id = ?", (file_id,))
-                return [pair for pair in arcs]
+                query = "select fromno, tono from arc where file_id = ?"
+                data = [file_id]
+                if context is not None:
+                    query += " and context_id = ?"
+                    data += [self._context_id(context)]
+                arcs = con.execute(query, data)
+                return list(arcs)
 
     def run_infos(self):
         return []   # TODO

tests/test_context.py

@@ -3,7 +3,10 @@
 
 """Tests for context support."""
 
+import os.path
+
 import coverage
+from coverage.data import CoverageData, combine_parallel_data
 
 from tests.coveragetest import CoverageTest
 
@@ -28,3 +31,59 @@ def test_global_context(self):
         self.start_import_stop(cov, "main")
         data = cov.get_data()
         self.assertCountEqual(data.measured_contexts(), ["gooey"])
+
+    def run_red_blue(self, **options):
+        self.make_file("red.py", """\
+            a = 1
+            if a > 2:
+                a = 3
+            assert a == 1
+            """)
+        red_cov = coverage.Coverage(context="red", data_suffix="r", source=["."], **options)
+        self.start_import_stop(red_cov, "red")
+        red_cov.save()
+
+        self.make_file("blue.py", """\
+            b = 1
+            if b > 2:
+                b = 3
+            assert b == 1
+            """)
+        blue_cov = coverage.Coverage(context="blue", data_suffix="b", source=["."], **options)
+        self.start_import_stop(blue_cov, "blue")
+        blue_cov.save()
+
+    def test_combining_line_contexts(self):
+        self.run_red_blue()
+        combined = CoverageData()
+        combine_parallel_data(combined)
+
+        self.assertEqual(combined.measured_contexts(), {'red', 'blue'})
+
+        full_names = {os.path.basename(f): f for f in combined.measured_files()}
+        self.assertCountEqual(full_names, ['red.py', 'blue.py'])
+
+        self.assertEqual(combined.lines(full_names['red.py'], context='red'), [1, 2, 4])
+        self.assertEqual(combined.lines(full_names['red.py'], context='blue'), [])
+        self.assertEqual(combined.lines(full_names['blue.py'], context='red'), [])
+        self.assertEqual(combined.lines(full_names['blue.py'], context='blue'), [1, 2, 4])
+
+    def test_combining_arc_contexts(self):
+        self.run_red_blue(branch=True)
+        combined = CoverageData()
+        combine_parallel_data(combined)
+
+        self.assertEqual(combined.measured_contexts(), {'red', 'blue'})
+
+        full_names = {os.path.basename(f): f for f in combined.measured_files()}
+        self.assertCountEqual(full_names, ['red.py', 'blue.py'])
+
+        self.assertEqual(combined.lines(full_names['red.py'], context='red'), [1, 2, 4])
+        self.assertEqual(combined.lines(full_names['red.py'], context='blue'), [])
+        self.assertEqual(combined.lines(full_names['blue.py'], context='red'), [])
+        self.assertEqual(combined.lines(full_names['blue.py'], context='blue'), [1, 2, 4])
+
+        self.assertEqual(combined.arcs(full_names['red.py'], context='red'), [(-1, 1), (1, 2), (2, 4), (4, -1)])
+        self.assertEqual(combined.arcs(full_names['red.py'], context='blue'), [])
+        self.assertEqual(combined.arcs(full_names['blue.py'], context='red'), [])
+        self.assertEqual(combined.arcs(full_names['blue.py'], context='blue'), [(-1, 1), (1, 2), (2, 4), (4, -1)])