Setting query context has to be done with a separate method call

5 files changed, 112 insertions, 121 deletions

diff --git a/coverage/sqldata.py b/coverage/sqldata.py
index 06ccff52..95af6072 100644
--- a/coverage/sqldata.py
+++ b/coverage/sqldata.py
@@ -192,7 +192,6 @@ class CoverageData(SimpleReprMixin):
 
         self._current_context = None
         self._current_context_id = None
-        self._query_contexts = None
         self._query_context_ids = None
 
     def _choose_filename(self):
@@ -714,32 +713,27 @@ class CoverageData(SimpleReprMixin):
                 return row[0] or ""
             return ""   # File was measured, but no tracer associated.
 
+    def set_query_context(self, context):
+        self._start_using()
+        with self._connect() as con:
+            cur = con.execute("select id from context where context = ?", (context,))
+            self._query_context_ids = [row[0] for row in cur.fetchall()]
 
     def set_query_contexts(self, contexts):
         """Set query contexts for future `lines`, `arcs` etc. calls."""
+        self._start_using()
         if contexts:
-            self._query_context_ids = self._get_query_context_ids(contexts)
-        else:
-            self._query_context_ids = None
-        self._query_contexts = contexts
-
-    def _get_query_context_ids(self, contexts=None):
-        if contexts is not None:
-            if not contexts:
-                return None
-            self._start_using()
             with self._connect() as con:
                 context_clause = ' or '.join(['context glob ?'] * len(contexts))
                 cur = con.execute("select id from context where " + context_clause, contexts)
-                return [row[0] for row in cur.fetchall()]
-        elif self._query_contexts:
-            return self._query_context_ids
-        return None
+                self._query_context_ids = [row[0] for row in cur.fetchall()]
+        else:
+            self._query_context_ids = None
 
-    def lines(self, filename, contexts=None):
+    def lines(self, filename):
         self._start_using()
         if self.has_arcs():
-            arcs = self.arcs(filename, contexts=contexts)
+            arcs = self.arcs(filename)
             if arcs is not None:
                 all_lines = itertools.chain.from_iterable(arcs)
                 return list(set(l for l in all_lines if l > 0))
@@ -751,18 +745,17 @@ class CoverageData(SimpleReprMixin):
             else:
                 query = "select numbits from line_bits where file_id = ?"
                 data = [file_id]
-                context_ids = self._get_query_context_ids(contexts)
-                if context_ids is not None:
-                    ids_array = ', '.join('?'*len(context_ids))
+                if self._query_context_ids is not None:
+                    ids_array = ', '.join('?' * len(self._query_context_ids))
                     query += " and context_id in (" + ids_array + ")"
-                    data += context_ids
+                    data += self._query_context_ids
                 bitmaps = list(con.execute(query, data))
                 nums = set()
                 for row in bitmaps:
                     nums.update(numbits_to_nums(row[0]))
                 return sorted(nums)
 
-    def arcs(self, filename, contexts=None):
+    def arcs(self, filename):
         self._start_using()
         with self._connect() as con:
             file_id = self._file_id(filename)
@@ -771,11 +764,10 @@ class CoverageData(SimpleReprMixin):
             else:
                 query = "select distinct fromno, tono from arc where file_id = ?"
                 data = [file_id]
-                context_ids = self._get_query_context_ids(contexts)
-                if context_ids is not None:
-                    ids_array = ', '.join('?'*len(context_ids))
+                if self._query_context_ids is not None:
+                    ids_array = ', '.join('?' * len(self._query_context_ids))
                     query += " and context_id in (" + ids_array + ")"
-                    data += context_ids
+                    data += self._query_context_ids
                 arcs = con.execute(query, data)
                 return list(arcs)
 
@@ -793,11 +785,10 @@ class CoverageData(SimpleReprMixin):
                     "where arc.file_id = ? and arc.context_id = context.id"
                 )
                 data = [file_id]
-                context_ids = self._get_query_context_ids()
-                if context_ids is not None:
-                    ids_array = ', '.join('?'*len(context_ids))
+                if self._query_context_ids is not None:
+                    ids_array = ', '.join('?' * len(self._query_context_ids))
                     query += " and arc.context_id in (" + ids_array + ")"
-                    data += context_ids
+                    data += self._query_context_ids
                 for fromno, tono, context in con.execute(query, data):
                     if context not in lineno_contexts_map[fromno]:
                         lineno_contexts_map[fromno].append(context)
@@ -810,11 +801,10 @@ class CoverageData(SimpleReprMixin):
                     "and file_id = ?"
                     )
                 data = [file_id]
-                context_ids = self._get_query_context_ids()
-                if context_ids is not None:
-                    ids_array = ', '.join('?'*len(context_ids))
+                if self._query_context_ids is not None:
+                    ids_array = ', '.join('?' * len(self._query_context_ids))
                     query += " and l.context_id in (" + ids_array + ")"
-                    data += context_ids
+                    data += self._query_context_ids
                 for numbits, context in con.execute(query, data):
                     for lineno in numbits_to_nums(numbits):
                         lineno_contexts_map[lineno].append(context)
diff --git a/tests/test_api.py b/tests/test_api.py
index 4f248441..b2143548 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -539,10 +539,10 @@ class ApiTest(CoverageTest):
         filenames = self.get_measured_filenames(data)
         suite_filename = filenames['testsuite.py']
 
-        self.assertEqual(
-            [2, 8], data.lines(suite_filename, contexts=["multiply_six"]))
-        self.assertEqual(
-            [2, 5], data.lines(suite_filename, contexts=["multiply_zero"]))
+        data.set_query_context("multiply_six")
+        self.assertEqual([2, 8], data.lines(suite_filename))
+        data.set_query_context("multiply_zero")
+        self.assertEqual([2, 5], data.lines(suite_filename))
 
     def test_switch_context_with_static(self):
         # This test simulates a coverage-aware test runner,
@@ -579,10 +579,10 @@ class ApiTest(CoverageTest):
         filenames = self.get_measured_filenames(data)
         suite_filename = filenames['testsuite.py']
 
-        self.assertEqual(
-            [2, 8], data.lines(suite_filename, contexts=["mysuite|multiply_six"]))
-        self.assertEqual(
-            [2, 5], data.lines(suite_filename, contexts=["mysuite|multiply_zero"]))
+        data.set_query_context("mysuite|multiply_six")
+        self.assertEqual([2, 8], data.lines(suite_filename))
+        data.set_query_context("mysuite|multiply_zero")
+        self.assertEqual([2, 5], data.lines(suite_filename))
 
     def test_switch_context_unstarted(self):
         # Coverage must be started to switch context
diff --git a/tests/test_context.py b/tests/test_context.py
index eb8ee9c5..72739ad0 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -72,10 +72,14 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            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 assert_combined_lines(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.lines(filename), lines)
+
+            assert_combined_lines(fred, 'red', self.LINES)
+            assert_combined_lines(fred, 'blue', [])
+            assert_combined_lines(fblue, 'red', [])
+            assert_combined_lines(fblue, 'blue', self.LINES)
 
     def test_combining_arc_contexts(self):
         red_data, blue_data = self.run_red_blue(branch=True)
@@ -92,15 +96,23 @@ class StaticContextTest(CoverageTest):
             fred = full_names['red.py']
             fblue = full_names['blue.py']
 
-            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 assert_combined_lines(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.lines(filename), lines)
+
+            assert_combined_lines(fred, 'red', self.LINES)
+            assert_combined_lines(fred, 'blue', [])
+            assert_combined_lines(fblue, 'red', [])
+            assert_combined_lines(fblue, 'blue', self.LINES)
+
+            def assert_combined_arcs(filename, context, lines):
+                combined.set_query_context(context)
+                self.assertEqual(combined.arcs(filename), lines)
 
-            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)
+            assert_combined_arcs(fred, 'red', self.ARCS)
+            assert_combined_arcs(fred, 'blue', [])
+            assert_combined_arcs(fblue, 'red', [])
+            assert_combined_arcs(fblue, 'blue', self.ARCS)
 
 
 class DynamicContextTest(CoverageTest):
@@ -145,12 +157,14 @@ class DynamicContextTest(CoverageTest):
         self.assertCountEqual(
             data.measured_contexts(),
             ["", "two_tests.test_one", "two_tests.test_two"])
-        self.assertCountEqual(data.lines(fname, [""]), self.OUTER_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["two_tests.test_one"]),
-            self.TEST_ONE_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["two_tests.test_two"]), self.TEST_TWO_LINES)
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertCountEqual(lines, data.lines(fname))
+
+        assert_context_lines("", self.OUTER_LINES)
+        assert_context_lines("two_tests.test_one", self.TEST_ONE_LINES)
+        assert_context_lines("two_tests.test_two", self.TEST_TWO_LINES)
 
     def test_static_and_dynamic(self):
         self.make_file("two_tests.py", self.SOURCE)
@@ -164,12 +178,14 @@ class DynamicContextTest(CoverageTest):
         self.assertCountEqual(
             data.measured_contexts(),
             ["stat", "stat|two_tests.test_one", "stat|two_tests.test_two"])
-        self.assertCountEqual(
-            data.lines(fname, ["stat"]), self.OUTER_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["stat|two_tests.test_one"]), self.TEST_ONE_LINES)
-        self.assertCountEqual(
-            data.lines(fname, ["stat|two_tests.test_two"]), self.TEST_TWO_LINES)
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertCountEqual(lines, data.lines(fname))
+
+        assert_context_lines("stat", self.OUTER_LINES)
+        assert_context_lines("stat|two_tests.test_one", self.TEST_ONE_LINES)
+        assert_context_lines("stat|two_tests.test_two", self.TEST_TWO_LINES)
 
 
 def get_qualname():
diff --git a/tests/test_data.py b/tests/test_data.py
index eb389e35..71e18bb9 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -201,8 +201,10 @@ class CoverageDataTest(DataTestHelpers, CoverageTest):
         covdata.set_context('test_a')
         covdata.add_lines(LINES_1)
         self.assertEqual(covdata.lines('a.py'), [1, 2])
-        self.assertEqual(covdata.lines('a.py', contexts=['test*']), [1, 2])
-        self.assertEqual(covdata.lines('a.py', contexts=['other*']), [])
+        covdata.set_query_contexts(['test*'])
+        self.assertEqual(covdata.lines('a.py'), [1, 2])
+        covdata.set_query_contexts(['other*'])
+        self.assertEqual(covdata.lines('a.py'), [])
 
     def test_contexts_by_lineno_with_lines(self):
         covdata = CoverageData()
@@ -250,11 +252,11 @@ class CoverageDataTest(DataTestHelpers, CoverageTest):
         covdata = CoverageData()
         covdata.set_context('test_x')
         covdata.add_arcs(ARCS_3)
-        self.assertEqual(
-            covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
-        self.assertEqual(covdata.arcs(
-            'x.py', contexts=['test*']), [(-1, 1), (1, 2), (2, 3), (3, -1)])
-        self.assertEqual(covdata.arcs('x.py', contexts=['other*']), [])
+        self.assertEqual(covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
+        covdata.set_query_contexts(['test*'])
+        self.assertEqual(covdata.arcs('x.py'), [(-1, 1), (1, 2), (2, 3), (3, -1)])
+        covdata.set_query_contexts(['other*'])
+        self.assertEqual(covdata.arcs('x.py'), [])
 
     def test_contexts_by_lineno_with_arcs(self):
         covdata = CoverageData()
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index eb51c1b9..22796dc3 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -990,18 +990,12 @@ class DynamicContextPluginTest(CoverageTest):
             ['', 'doctest:HTML_TAG', 'test:HTML_TAG', 'test:RENDERERS'],
             sorted(data.measured_contexts()),
         )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["test:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["test:RENDERERS"]),
-        )
+        data.set_query_context("doctest:HTML_TAG")
+        self.assertEqual([2], data.lines(filenames['rendering.py']))
+        data.set_query_context("test:HTML_TAG")
+        self.assertEqual([2], data.lines(filenames['rendering.py']))
+        data.set_query_context("test:RENDERERS")
+        self.assertEqual([2, 5, 8, 11], data.lines(filenames['rendering.py']))
 
     def test_static_context(self):
         self.make_plugin_capitalized_testnames('plugin_tests.py')
@@ -1041,22 +1035,21 @@ class DynamicContextPluginTest(CoverageTest):
         # labeled by plugin_tests.
         data = cov.get_data()
         filenames = self.get_measured_filenames(data)
-        self.assertEqual(
-            ['', 'doctest:HTML_TAG', 'testsuite.test_html_tag', 'testsuite.test_renderers'],
-            sorted(data.measured_contexts()),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["testsuite.test_html_tag"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["testsuite.test_renderers"]),
-        )
+        expected = [
+            '',
+            'doctest:HTML_TAG',
+            'testsuite.test_html_tag',
+            'testsuite.test_renderers',
+        ]
+        self.assertEqual(expected, sorted(data.measured_contexts()))
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertEqual(lines, data.lines(filenames['rendering.py']))
+
+        assert_context_lines("doctest:HTML_TAG", [2])
+        assert_context_lines("testsuite.test_html_tag", [2])
+        assert_context_lines("testsuite.test_renderers", [2, 5, 8, 11])
 
     def test_multiple_plugins(self):
         self.make_plugin_capitalized_testnames('plugin_tests.py')
@@ -1087,23 +1080,13 @@ class DynamicContextPluginTest(CoverageTest):
             'test:RENDERERS',
         ]
         self.assertEqual(expected, sorted(data.measured_contexts()))
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["test:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5, 8, 11],
-            data.lines(filenames['rendering.py'], contexts=["test:RENDERERS"]),
-        )
-        self.assertEqual(
-            [2],
-            data.lines(filenames['rendering.py'], contexts=["doctest:HTML_TAG"]),
-        )
-        self.assertEqual(
-            [2, 5],
-            data.lines(filenames['rendering.py'], contexts=["renderer:paragraph"]),
-        )
-        self.assertEqual(
-            [2, 8],
-            data.lines(filenames['rendering.py'], contexts=["renderer:span"]),
-        )
+
+        def assert_context_lines(context, lines):
+            data.set_query_context(context)
+            self.assertEqual(lines, data.lines(filenames['rendering.py']))
+
+        assert_context_lines("test:HTML_TAG", [2])
+        assert_context_lines("test:RENDERERS", [2, 5, 8, 11])
+        assert_context_lines("doctest:HTML_TAG", [2])
+        assert_context_lines("renderer:paragraph", [2, 5])
+        assert_context_lines("renderer:span", [2, 8])