Properly round-trip the trace function even when not measuring coverage. #575

4 files changed, 48 insertions, 30 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 76a1172..2f57279 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -10,6 +10,11 @@ Change history for Coverage.py
 Unreleased
 ----------
 
+- Code that uses ``sys.settrace(sys.gettrace())`` in a file that wasn't being
+  coverage-measured would prevent correct coverage measurement in following
+  code. An example of this was running doctests programmatically, as described
+  in `issue 575`_. This is now fixed.
+
 - Errors printed by the ``coverage`` command now go to stderr instead of
   stdout.
 
@@ -17,6 +22,7 @@ Unreleased
   fail under Python 2, as reported in `issue 573`_. This is now fixed.
 
 .. _issue 573: https://bitbucket.org/ned/coveragepy/issues/573/cant-generate-xml-report-if-some-source
+.. _issue 575: https://bitbucket.org/ned/coveragepy/issues/575/running-doctest-prevents-complete-coverage
 
 
 Version 4.4b1 --- 2017-04-04
diff --git a/coverage/ctracer/tracer.c b/coverage/ctracer/tracer.c
index ee112d8..095df11 100644
--- a/coverage/ctracer/tracer.c
+++ b/coverage/ctracer/tracer.c
@@ -529,9 +529,6 @@ CTracer_handle_call(CTracer *self, PyFrameObject *frame)
         self->pcur_entry->file_data = file_data;
         self->pcur_entry->file_tracer = file_tracer;
 
-        /* Make the frame right in case settrace(gettrace()) happens. */
-        Py_INCREF(self);
-        frame->f_trace = (PyObject*)self;
         SHOWLOG(self->pdata_stack->depth, frame->f_lineno, filename, "traced");
     }
     else {
@@ -543,6 +540,10 @@ CTracer_handle_call(CTracer *self, PyFrameObject *frame)
 
     self->pcur_entry->disposition = disposition;
 
+    /* Make the frame right in case settrace(gettrace()) happens. */
+    Py_INCREF(self);
+    My_XSETREF(frame->f_trace, (PyObject*)self);
+
     /* A call event is really a "start frame" event, and can happen for
      * re-entering a generator also.  f_lasti is -1 for a true call, and a
      * real byte offset for a generator re-entry.
diff --git a/coverage/ctracer/util.h b/coverage/ctracer/util.h
index e192f97..f0c302c 100644
--- a/coverage/ctracer/util.h
+++ b/coverage/ctracer/util.h
@@ -44,6 +44,14 @@
 
 #endif /* Py3k */
 
+// Undocumented, and not in 2.6, so our own copy of it.
+#define My_XSETREF(op, op2)                     \
+    do {                                        \
+        PyObject *_py_tmp = (PyObject *)(op);   \
+        (op) = (op2);                           \
+        Py_XDECREF(_py_tmp);                    \
+    } while (0)
+
 /* The values returned to indicate ok or error. */
 #define RET_OK      0
 #define RET_ERROR   -1
diff --git a/tests/test_oddball.py b/tests/test_oddball.py
index 91acc0c..b54f4ef 100644
--- a/tests/test_oddball.py
+++ b/tests/test_oddball.py
@@ -432,35 +432,38 @@ class DoctestTest(CoverageTest):
 
 class GettraceTest(CoverageTest):
     """Tests that we work properly with `sys.gettrace()`."""
-    def test_round_trip(self):
-        self.check_coverage('''\
-            import sys
-            def foo(n):
-                return 3*n
-            def bar(n):
-                return 5*n
-            a = foo(6)
-            sys.settrace(sys.gettrace())
-            a = bar(8)
-            ''',
-            [1, 2, 3, 4, 5, 6, 7, 8], "")
-
-    def test_multi_layers(self):
-        self.check_coverage('''\
+    def test_round_trip_in_untraced_function(self):
+        # https://bitbucket.org/ned/coveragepy/issues/575/running-doctest-prevents-complete-coverage
+        self.make_file("main.py", """import sample""")
+        self.make_file("sample.py", """\
+            from swap import swap_it
+            def doit():
+                print(3)
+                swap_it()
+                print(5)
+            def doit_soon():
+                print(7)
+                doit()
+                print(9)
+            print(10)
+            doit_soon()
+            print(12)
+            """)
+        self.make_file("swap.py", """\
             import sys
-            def level1():
-                a = 3
-                level2()
-                b = 5
-            def level2():
-                c = 7
+            def swap_it():
                 sys.settrace(sys.gettrace())
-                d = 9
-            e = 10
-            level1()
-            f = 12
-            ''',
-            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], "")
+            """)
+
+        # Use --source=sample to prevent measurement of swap.py.
+        cov = coverage.Coverage(source=["sample"])
+        self.start_import_stop(cov, "main")
+
+        self.assertEqual(self.stdout(), "10\n7\n3\n5\n9\n12\n")
+
+        _, statements, missing, _ = cov.analysis("sample.py")
+        self.assertEqual(statements, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+        self.assertEqual(missing, [])
 
     def test_setting_new_trace_function(self):
         # https://bitbucket.org/ned/coveragepy/issues/436/disabled-coverage-ctracer-may-rise-from