Move disable_plugin to Python

3 files changed, 21 insertions, 47 deletions

diff --git a/coverage/collector.py b/coverage/collector.py
index a042357f..3545ab1e 100644
--- a/coverage/collector.py
+++ b/coverage/collector.py
@@ -256,6 +256,8 @@ class Collector(object):
         if hasattr(tracer, 'should_start_context'):
             tracer.should_start_context = self.should_start_context
             tracer.switch_context = self.switch_context
+        if hasattr(tracer, 'disable_plugin'):
+            tracer.disable_plugin = self.disable_plugin
 
         fn = tracer.start()
         self.tracers.append(tracer)
@@ -381,6 +383,15 @@ class Collector(object):
             context = new_context
         self.covdata.set_context(context)
 
+    def disable_plugin(self, disposition):
+        """Disable the plugin mentioned in `disposition`."""
+        file_tracer = disposition.file_tracer
+        plugin = file_tracer._coverage_plugin
+        plugin_name = plugin._coverage_plugin_name
+        self.warn("Disabling plug-in {!r} due to previous exception".format(plugin_name))
+        plugin._coverage_enabled = False
+        disposition.trace = False
+
     def cached_mapped_file(self, filename):
         """A locally cached version of file names mapped through file_mapper."""
         key = (type(filename), filename)
diff --git a/coverage/ctracer/tracer.c b/coverage/ctracer/tracer.c
index 50808642..04552352 100644
--- a/coverage/ctracer/tracer.c
+++ b/coverage/ctracer/tracer.c
@@ -103,6 +103,7 @@ CTracer_dealloc(CTracer *self)
     Py_XDECREF(self->should_start_context);
     Py_XDECREF(self->switch_context);
     Py_XDECREF(self->context);
+    Py_XDECREF(self->disable_plugin);
 
     DataStack_dealloc(&self->stats, &self->data_stack);
     if (self->data_stacks) {
@@ -570,52 +571,17 @@ error:
 static void
 CTracer_disable_plugin(CTracer *self, PyObject * disposition)
 {
-    PyObject * file_tracer = NULL;
-    PyObject * plugin = NULL;
-    PyObject * plugin_name = NULL;
-    PyObject * msg = NULL;
-    PyObject * ignored = NULL;
-
+    PyObject * ret;
     PyErr_Print();
 
-    file_tracer = PyObject_GetAttr(disposition, str_file_tracer);
-    if (file_tracer == NULL) {
-        goto error;
-    }
-    if (file_tracer == Py_None) {
-        /* This shouldn't happen... */
-        goto ok;
-    }
-    plugin = PyObject_GetAttr(file_tracer, str__coverage_plugin);
-    if (plugin == NULL) {
-        goto error;
-    }
-    plugin_name = PyObject_GetAttr(plugin, str__coverage_plugin_name);
-    if (plugin_name == NULL) {
-        goto error;
-    }
-    msg = MyText_FromFormat(
-        "Disabling plug-in '%s' due to previous exception",
-        MyText_AsString(plugin_name)
-        );
-    if (msg == NULL) {
-        goto error;
-    }
     STATS( self->stats.pycalls++; )
-    ignored = PyObject_CallFunctionObjArgs(self->warn, msg, NULL);
-    if (ignored == NULL) {
-        goto error;
-    }
-
-    /* Disable the plugin for future files, and stop tracing this file. */
-    if (PyObject_SetAttr(plugin, str__coverage_enabled, Py_False) < 0) {
-        goto error;
-    }
-    if (PyObject_SetAttr(disposition, str_trace, Py_False) < 0) {
+    ret = PyObject_CallFunctionObjArgs(self->disable_plugin, disposition, NULL);
+    if (ret == NULL) {
         goto error;
     }
+    Py_DECREF(ret);
 
-    goto ok;
+    return;
 
 error:
     /* This function doesn't return a status, so if an error happens, print it,
@@ -623,13 +589,6 @@ error:
     /* PySys_WriteStderr is nicer, but is not in the public API. */
     fprintf(stderr, "Error occurred while disabling plug-in:\n");
     PyErr_Print();
-
-ok:
-    Py_XDECREF(file_tracer);
-    Py_XDECREF(plugin);
-    Py_XDECREF(plugin_name);
-    Py_XDECREF(msg);
-    Py_XDECREF(ignored);
 }
 
 
@@ -1121,6 +1080,9 @@ CTracer_members[] = {
     { "switch_context",     T_OBJECT, offsetof(CTracer, switch_context), 0,
             PyDoc_STR("Function for switching to a new context.") },
 
+    { "disable_plugin",     T_OBJECT, offsetof(CTracer, disable_plugin), 0,
+            PyDoc_STR("Function for disabling a plugin.") },
+
     { NULL }
 };
 
diff --git a/coverage/ctracer/tracer.h b/coverage/ctracer/tracer.h
index a83742dd..8994a9e3 100644
--- a/coverage/ctracer/tracer.h
+++ b/coverage/ctracer/tracer.h
@@ -27,6 +27,7 @@ typedef struct CTracer {
     PyObject * trace_arcs;
     PyObject * should_start_context;
     PyObject * switch_context;
+    PyObject * disable_plugin;
 
     /* Has the tracer been started? */
     BOOL started;