Support for multiple, glob-based contexts in result queries.

3 files changed, 61 insertions, 25 deletions

diff --git a/coverage/data.py b/coverage/data.py
index 59bb073b..e393345a 100644
--- a/coverage/data.py
+++ b/coverage/data.py
@@ -175,6 +175,13 @@ class CoverageJsonData(object):
     ## Reading data
     ##
 
+    def set_query_contexts(self, contexts=None):
+        """Set the query contexts.
+
+        No-op, since contexts are not supported for this data format.
+        """
+        pass
+
     def has_arcs(self):
         """Does this data have arcs?
 
@@ -186,7 +193,7 @@ class CoverageJsonData(object):
         """
         return self._has_arcs()
 
-    def lines(self, filename):
+    def lines(self, filename, contexts=None):
         """Get the list of lines executed for a file.
 
         If the file was not measured, returns None.  A file might be measured,
@@ -195,6 +202,8 @@ class CoverageJsonData(object):
         If the file was executed, returns a list of integers, the line numbers
         executed in the file. The list is in no particular order.
 
+        `contexts` is ignored, since contexts are not supported for this data
+        format.
         """
         if self._arcs is not None:
             arcs = self._arcs.get(filename)
@@ -205,7 +214,7 @@ class CoverageJsonData(object):
             return self._lines.get(filename)
         return None
 
-    def arcs(self, filename):
+    def arcs(self, filename, contexts=None):
         """Get the list of arcs executed for a file.
 
         If the file was not measured, returns None.  A file might be measured,
@@ -221,6 +230,8 @@ class CoverageJsonData(object):
         If the ending ling number is -N, it's an exit from the code object that
         starts at line N.
 
+        `contexts` is ignored, since contexts are not supported for this data
+        format.
         """
         if self._arcs is not None:
             if filename in self._arcs:
diff --git a/coverage/sqldata.py b/coverage/sqldata.py
index f8961fa6..5ba393d4 100644
--- a/coverage/sqldata.py
+++ b/coverage/sqldata.py
@@ -546,7 +546,30 @@ class CoverageSqliteData(SimpleReprMixin):
                 return row[0] or ""
             return ""   # File was measured, but no tracer associated.
 
-    def lines(self, filename, context=None):
+
+    def set_query_contexts(self, contexts=None):
+        """Set query contexts for future `lines`, `arcs` etc. calls."""
+        self._query_context_ids = self._get_query_context_ids(contexts) \
+            if contexts is not None else None
+        self._query_contexts = contexts
+
+    def _get_query_context_ids(self, contexts=None):
+        if contexts is not None:
+            if not len(contexts):
+                return None
+            self._start_using()
+            with self._connect() as con:
+                # Context entries can be globs, so convert '*' with '%'.
+                context_selectors = [context.replace('*', '%') for context in contexts]
+                context_clause = ' or '.join(['contenxt like ?']*len(contexts))
+                con.execute(
+                    "select id from context where " + context_clause, context_selectors)
+                return [row[0] for row in con.fetchall()]
+        elif self._query_contexts is not None:
+            return self._query_context_ids
+        return None
+
+    def lines(self, filename, contexts=None):
         self._start_using()
         if self.has_arcs():
             arcs = self.arcs(filename, context=context)
@@ -561,13 +584,14 @@ class CoverageSqliteData(SimpleReprMixin):
             else:
                 query = "select distinct lineno from line where file_id = ?"
                 data = [file_id]
-                if context is not None:
-                    query += " and context_id = ?"
-                    data += [self._context_id(context)]
+                context_ids = self._get_query_context_ids(contexts)
+                if context_ids is not None:
+                    query += " and context_id IN ?"
+                    data += [context_ids]
                 linenos = con.execute(query, data)
                 return [lineno for lineno, in linenos]
 
-    def arcs(self, filename, context=None):
+    def arcs(self, filename, contexts=None):
         self._start_using()
         with self._connect() as con:
             file_id = self._file_id(filename)
@@ -576,9 +600,10 @@ class CoverageSqliteData(SimpleReprMixin):
             else:
                 query = "select distinct fromno, tono from arc where file_id = ?"
                 data = [file_id]
-                if context is not None:
-                    query += " and context_id = ?"
-                    data += [self._context_id(context)]
+                context_ids = self._get_query_context_ids(contexts)
+                if context_ids is not None:
+                    query += " and context_id IN ?"
+                    data += [context_ids]
                 arcs = con.execute(query, data)
                 return list(arcs)
 
diff --git a/tests/test_context.py b/tests/test_context.py
index 29839c57..726070e3 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -77,10 +77,10 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            self.assertEqual(combined.lines(fred, context='red'), self.LINES)
-            self.assertEqual(combined.lines(fred, context='blue'), [])
-            self.assertEqual(combined.lines(fblue, context='red'), [])
-            self.assertEqual(combined.lines(fblue, context='blue'), self.LINES)
+            self.assertEqual(combined.lines(fred, contexts=['red']), self.LINES)
+            self.assertEqual(combined.lines(fred, contexts=['blue']), [])
+            self.assertEqual(combined.lines(fblue, contexts=['red']), [])
+            self.assertEqual(combined.lines(fblue, contexts=['blue']), self.LINES)
 
     def test_combining_arc_contexts(self):
         red_data, blue_data = self.run_red_blue(branch=True)
@@ -97,15 +97,15 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            self.assertEqual(combined.lines(fred, context='red'), self.LINES)
-            self.assertEqual(combined.lines(fred, context='blue'), [])
-            self.assertEqual(combined.lines(fblue, context='red'), [])
-            self.assertEqual(combined.lines(fblue, context='blue'), self.LINES)
+            self.assertEqual(combined.lines(fred, contexts=['red']), self.LINES)
+            self.assertEqual(combined.lines(fred, contexts=['blue']), [])
+            self.assertEqual(combined.lines(fblue, contexts=['red']), [])
+            self.assertEqual(combined.lines(fblue, contexts=['blue']), self.LINES)
 
-            self.assertEqual(combined.arcs(fred, context='red'), self.ARCS)
-            self.assertEqual(combined.arcs(fred, context='blue'), [])
-            self.assertEqual(combined.arcs(fblue, context='red'), [])
-            self.assertEqual(combined.arcs(fblue, context='blue'), self.ARCS)
+            self.assertEqual(combined.arcs(fred, contexts=['red']), self.ARCS)
+            self.assertEqual(combined.arcs(fred, contexts=['blue']), [])
+            self.assertEqual(combined.arcs(fblue, contexts=['red']), [])
+            self.assertEqual(combined.arcs(fblue, contexts=['blue']), self.ARCS)
 
 
 class DynamicContextTest(CoverageTest):
@@ -154,9 +154,9 @@ class DynamicContextTest(CoverageTest):
         full_names = {os.path.basename(f): f for f in data.measured_files()}
         fname = full_names["two_tests.py"]
         self.assertCountEqual(data.measured_contexts(), ["", "test_one", "test_two"])
-        self.assertCountEqual(data.lines(fname, ""), self.OUTER_LINES)
-        self.assertCountEqual(data.lines(fname, "test_one"), self.TEST_ONE_LINES)
-        self.assertCountEqual(data.lines(fname, "test_two"), self.TEST_TWO_LINES)
+        self.assertCountEqual(data.lines(fname, [""]), self.OUTER_LINES)
+        self.assertCountEqual(data.lines(fname, ["test_one"]), self.TEST_ONE_LINES)
+        self.assertCountEqual(data.lines(fname, ["test_two"]), self.TEST_TWO_LINES)
 
     def test_static_and_dynamic(self):
         self.make_file("two_tests.py", self.SOURCE)