Fix settrace(py_func). #436.

4 files changed, 85 insertions, 10 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 60fb44f7..3653e36b 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -14,11 +14,15 @@ Version 4.0.2
 - Files or directories with non-ASCII characters are now handled properly,
   fixing `issue 432`_.
 
+- Setting a trace function with sys.settrace was broken by a change in 4.0.1,
+  as reported in `issue 436`_.  This is now fixed.
+
 - Officially support PyPy 4.0, which required no changes, just updates to the
   docs.
 
 .. _issue 431: https://bitbucket.org/ned/coveragepy/issues/431/couldnt-parse-python-file-with-cp1252
 .. _issue 432: https://bitbucket.org/ned/coveragepy/issues/432/path-with-unicode-characters-various
+.. _issue 436: https://bitbucket.org/ned/coveragepy/issues/436/disabled-coverage-ctracer-may-rise-from
 
 
 Version 4.0.1 --- 13 October 2015
diff --git a/coverage/ctracer/tracer.c b/coverage/ctracer/tracer.c
index e5f39d09..79bebac5 100644
--- a/coverage/ctracer/tracer.c
+++ b/coverage/ctracer/tracer.c
@@ -768,7 +768,7 @@ CTracer_trace(CTracer *self, PyFrameObject *frame, int what, PyObject *arg_unuse
     #endif
 
     #if WHAT_LOG
-    if (what <= sizeof(what_sym)/sizeof(const char *)) {
+    if (what <= (int)(sizeof(what_sym)/sizeof(const char *))) {
         ascii = MyText_AS_BYTES(frame->f_code->co_filename);
         printf("trace: %s @ %s %d\n", what_sym[what], MyBytes_AS_STRING(ascii), frame->f_lineno);
         Py_DECREF(ascii);
@@ -859,6 +859,7 @@ CTracer_call(CTracer *self, PyObject *args, PyObject *kwds)
     int what;
     int orig_lineno;
     PyObject *ret = NULL;
+    PyObject * ascii = NULL;
 
     static char *what_names[] = {
         "call", "exception", "line", "return",
@@ -866,10 +867,6 @@ CTracer_call(CTracer *self, PyObject *args, PyObject *kwds)
         NULL
         };
 
-    #if WHAT_LOG
-    printf("pytrace\n");
-    #endif
-
     static char *kwlist[] = {"frame", "event", "arg", "lineno", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!O!O|i:Tracer_call", kwlist,
@@ -880,7 +877,7 @@ CTracer_call(CTracer *self, PyObject *args, PyObject *kwds)
     /* In Python, the what argument is a string, we need to find an int
        for the C function. */
     for (what = 0; what_names[what]; what++) {
-        PyObject *ascii = MyText_AS_BYTES(what_str);
+        ascii = MyText_AS_BYTES(what_str);
         int should_break = !strcmp(MyBytes_AS_STRING(ascii), what_names[what]);
         Py_DECREF(ascii);
         if (should_break) {
@@ -888,6 +885,12 @@ CTracer_call(CTracer *self, PyObject *args, PyObject *kwds)
         }
     }
 
+    #if WHAT_LOG
+    ascii = MyText_AS_BYTES(frame->f_code->co_filename);
+    printf("pytrace: %s @ %s %d\n", what_sym[what], MyBytes_AS_STRING(ascii), frame->f_lineno);
+    Py_DECREF(ascii);
+    #endif
+
     /* Save off the frame's lineno, and use the forced one, if provided. */
     orig_lineno = frame->f_lineno;
     if (lineno > 0) {
@@ -904,8 +907,30 @@ CTracer_call(CTracer *self, PyObject *args, PyObject *kwds)
     frame->f_lineno = orig_lineno;
 
     /* For better speed, install ourselves the C way so that future calls go
-       directly to CTracer_trace, without this intermediate function. */
-    PyEval_SetTrace((Py_tracefunc)CTracer_trace, (PyObject*)self);
+       directly to CTracer_trace, without this intermediate function.
+       
+       Only do this if this is a CALL event, since new trace functions only
+       take effect then.  If we don't condition it on CALL, then we'll clobber
+       the new trace function before it has a chance to get called.  To
+       understand why, there are three internal values to track: frame.f_trace,
+       c_tracefunc, and c_traceobj.  They are explained here:
+       http://nedbatchelder.com/text/trace-function.html
+
+       Without the conditional on PyTrace_CALL, this is what happens:
+
+            def func():                 #   f_trace         c_tracefunc     c_traceobj
+                                        #   --------------  --------------  --------------
+                                        #   CTracer         CTracer.trace   CTracer
+                sys.settrace(my_func)   
+                                        #   CTracer         trampoline      my_func
+                        # Now Python calls trampoline(CTracer), which calls this function
+                        # which calls PyEval_SetTrace below, setting us as the tracer again:
+                                        #   CTracer         CTracer.trace   CTracer
+                        # and it's as if the settrace never happened.
+        */
+    if (what == PyTrace_CALL) {
+        PyEval_SetTrace((Py_tracefunc)CTracer_trace, (PyObject*)self);
+    }
 
 done:
     return ret;
diff --git a/tests/coveragetest.py b/tests/coveragetest.py
index 0e9076cc..5042c98d 100644
--- a/tests/coveragetest.py
+++ b/tests/coveragetest.py
@@ -48,6 +48,11 @@ class CoverageTest(
     def setUp(self):
         super(CoverageTest, self).setUp()
 
+        # Attributes for getting info about what happened.
+        self.last_command_status = None
+        self.last_command_output = None
+        self.last_module_name = None
+
         if _TEST_NAME_FILE:                                 # pragma: debugging
             with open(_TEST_NAME_FILE, "w") as f:
                 f.write("%s_%s" % (
@@ -104,8 +109,9 @@ class CoverageTest(
         return mod
 
     def get_module_name(self):
-        """Return the module name to use for this test run."""
-        return 'coverage_test_' + str(random.random())[2:]
+        """Return a random module name to use for this test run."""
+        self.last_module_name = 'coverage_test_' + str(random.random())[2:]
+        return self.last_module_name
 
     # Map chars to numbers for arcz_to_arcs
     _arcz_map = {'.': -1}
diff --git a/tests/test_oddball.py b/tests/test_oddball.py
index 8f9e9707..87c65b0e 100644
--- a/tests/test_oddball.py
+++ b/tests/test_oddball.py
@@ -410,6 +410,46 @@ class GettraceTest(CoverageTest):
             ''',
             [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], "")
 
+    def test_setting_new_trace_function(self):
+        # https://bitbucket.org/ned/coveragepy/issues/436/disabled-coverage-ctracer-may-rise-from
+        self.check_coverage('''\
+            import sys
+
+            def tracer(frame, event, arg):
+                print("%s: %s @ %d" % (event, frame.f_code.co_filename, frame.f_lineno))
+                return tracer
+
+            def begin():
+                sys.settrace(tracer)
+
+            def collect():
+                t = sys.gettrace()
+                assert t is tracer, t
+
+            def test_unsets_trace():
+                begin()
+                collect()
+
+            old = sys.gettrace()
+            test_unsets_trace()
+            sys.settrace(old)
+            ''',
+            lines=[1, 3, 4, 5, 7, 8, 10, 11, 12, 14, 15, 16, 18, 19, 20],
+            missing="4-5, 11-12",
+        )
+
+        out = self.stdout().replace(self.last_module_name, "coverage_test")
+
+        self.assertEqual(
+            out,
+            (
+                "call: coverage_test.py @ 10\n"
+                "line: coverage_test.py @ 11\n"
+                "line: coverage_test.py @ 12\n"
+                "return: coverage_test.py @ 12\n"
+            ),
+        )
+
 
 class ExecTest(CoverageTest):
     """Tests of exec."""