Merged who-tests-what-170

6 files changed, 106 insertions, 7 deletions

diff --git a/coverage/collector.py b/coverage/collector.py
index 0a43d87c..aabf10b7 100644
--- a/coverage/collector.py
+++ b/coverage/collector.py
@@ -37,6 +37,11 @@ class FileDisposition(object):
     pass
 
 
+def should_start_context(frame):
+    fn_name = frame.f_code.co_name
+    if fn_name.startswith("test"):
+        return fn_name
+
 class Collector(object):
     """Collects trace data.
 
@@ -116,6 +121,10 @@ class Collector(object):
                 "Couldn't trace with concurrency=%s, the module isn't installed." % concurrency
             )
 
+        # Who-Tests-What is just a hack at the moment, so turn it on with an
+        # environment variable.
+        self.wtw = int(os.getenv('COVERAGE_WTW', 0))
+
         self.reset()
 
         if timid:
@@ -147,6 +156,10 @@ class Collector(object):
         # pairs as keys (if branch coverage).
         self.data = {}
 
+        # A dict mapping contexts to data dictionaries.
+        self.contexts = {}
+        self.contexts[None] = self.data
+
         # A dictionary mapping file names to file tracer plugin names that will
         # handle them.
         self.file_tracers = {}
@@ -206,6 +219,11 @@ class Collector(object):
             tracer.threading = self.threading
         if hasattr(tracer, 'check_include'):
             tracer.check_include = self.check_include
+        if self.wtw:
+            if hasattr(tracer, 'should_start_context'):
+                tracer.should_start_context = should_start_context
+            if hasattr(tracer, 'switch_context'):
+                tracer.switch_context = self.switch_context
 
         fn = tracer.start()
         self.tracers.append(tracer)
@@ -294,7 +312,7 @@ class Collector(object):
             if stats:
                 print("\nCoverage.py tracer stats:")
                 for k in sorted(stats.keys()):
-                    print("%16s: %s" % (k, stats[k]))
+                    print("%20s: %s" % (k, stats[k]))
         if self.threading:
             self.threading.settrace(None)
 
@@ -307,6 +325,11 @@ class Collector(object):
         else:
             self._start_tracer()
 
+    def switch_context(self, new_context):
+        data = self.contexts.setdefault(new_context, {})
+        for tracer in self.tracers:
+            tracer.data = data
+
     def save_data(self, covdata):
         """Save the collected data to a `CoverageData`.
 
@@ -323,4 +346,10 @@ class Collector(object):
             covdata.add_lines(abs_file_dict(self.data))
         covdata.add_file_tracers(abs_file_dict(self.file_tracers))
 
+        if self.wtw:
+            # Just a hack, so just hack it.
+            import pprint
+            with open("coverage_wtw.py", "w") as wtw_out:
+                pprint.pprint(self.contexts, wtw_out)
+
         self.reset()
diff --git a/coverage/ctracer/datastack.h b/coverage/ctracer/datastack.h
index 78f85f7e..13149798 100644
--- a/coverage/ctracer/datastack.h
+++ b/coverage/ctracer/datastack.h
@@ -27,6 +27,8 @@ typedef struct DataStackEntry {
         -1 means there was no previous line, as when entering a code object.
     */
     int last_line;
+
+    int started_context;
 } DataStackEntry;
 
 /* A data stack is a dynamically allocated vector of DataStackEntry's. */
diff --git a/coverage/ctracer/stats.h b/coverage/ctracer/stats.h
index ceba79bd..a72117cf 100644
--- a/coverage/ctracer/stats.h
+++ b/coverage/ctracer/stats.h
@@ -24,6 +24,7 @@ typedef struct Stats {
     unsigned int stack_reallocs;
     unsigned int errors;
     unsigned int pycalls;
+    unsigned int start_context_calls;
 #endif
 } Stats;
 
diff --git a/coverage/ctracer/tracer.c b/coverage/ctracer/tracer.c
index 681c9a97..02a613c0 100644
--- a/coverage/ctracer/tracer.c
+++ b/coverage/ctracer/tracer.c
@@ -73,6 +73,8 @@ CTracer_init(CTracer *self, PyObject *args_unused, PyObject *kwds_unused)
 
     self->cur_entry.last_line = -1;
 
+    self->context = Py_None;
+
     ret = RET_OK;
     goto ok;
 
@@ -99,6 +101,9 @@ CTracer_dealloc(CTracer *self)
     Py_XDECREF(self->data);
     Py_XDECREF(self->file_tracers);
     Py_XDECREF(self->should_trace_cache);
+    Py_XDECREF(self->should_start_context);
+    Py_XDECREF(self->switch_context);
+    Py_XDECREF(self->context);
 
     DataStack_dealloc(&self->stats, &self->data_stack);
     if (self->data_stacks) {
@@ -126,7 +131,7 @@ indent(int n)
     return spaces + strlen(spaces) - n*2;
 }
 
-static int logging = 0;
+static int logging = FALSE;
 /* Set these constants to be a file substring and line number to start logging. */
 static const char * start_file = "tests/views";
 static int start_line = 27;
@@ -338,6 +343,7 @@ CTracer_handle_call(CTracer *self, PyFrameObject *frame)
 
 
     STATS( self->stats.calls++; )
+
     /* Grow the stack. */
     if (CTracer_set_pdata_stack(self) < 0) {
         goto error;
@@ -349,6 +355,37 @@ CTracer_handle_call(CTracer *self, PyFrameObject *frame)
     /* Push the current state on the stack. */
     self->pdata_stack->stack[self->pdata_stack->depth] = self->cur_entry;
 
+    /* See if this frame begins a new context. */
+    if (self->should_start_context && self->context == Py_None) {
+        PyObject * context;
+        /* We're looking for our context, ask should_start_context if this is the start. */
+        STATS( self->stats.start_context_calls++; )
+        STATS( self->stats.pycalls++; )
+        context = PyObject_CallFunctionObjArgs(self->should_start_context, frame, NULL);
+        if (context == NULL) {
+            goto error;
+        }
+        if (context != Py_None) {
+            PyObject * val;
+            Py_DECREF(self->context);
+            self->context = context;
+            self->cur_entry.started_context = TRUE;
+            STATS( self->stats.pycalls++; )
+            val = PyObject_CallFunctionObjArgs(self->switch_context, context, NULL);
+            if (val == NULL) {
+                goto error;
+            }
+            Py_DECREF(val);
+        }
+        else {
+            Py_DECREF(context);
+            self->cur_entry.started_context = FALSE;
+        }
+    }
+    else {
+        self->cur_entry.started_context = FALSE;
+    }
+
     /* Check if we should trace this line. */
     filename = frame->f_code->co_filename;
     disposition = PyDict_GetItem(self->should_trace_cache, filename);
@@ -722,6 +759,22 @@ CTracer_handle_return(CTracer *self, PyFrameObject *frame)
             }
         }
 
+        /* If this frame started a context, then returning from it ends the context. */
+        if (self->cur_entry.started_context) {
+            PyObject * val;
+            Py_DECREF(self->context);
+            self->context = Py_None;
+            Py_INCREF(self->context);
+            STATS( self->stats.pycalls++; )
+
+            val = PyObject_CallFunctionObjArgs(self->switch_context, self->context, NULL);
+            if (val == NULL) {
+                goto error;
+            }
+            Py_DECREF(val);
+        }
+
+        /* Pop the stack. */
         SHOWLOG(self->pdata_stack->depth, frame->f_lineno, frame->f_code->co_filename, "return");
         self->cur_entry = self->pdata_stack->stack[self->pdata_stack->depth];
         self->pdata_stack->depth--;
@@ -781,7 +834,7 @@ CTracer_trace(CTracer *self, PyFrameObject *frame, int what, PyObject *arg_unuse
     #if TRACE_LOG
     ascii = MyText_AS_BYTES(frame->f_code->co_filename);
     if (strstr(MyBytes_AS_STRING(ascii), start_file) && frame->f_lineno == start_line) {
-        logging = 1;
+        logging = TRUE;
     }
     Py_DECREF(ascii);
     #endif
@@ -944,7 +997,7 @@ static PyObject *
 CTracer_start(CTracer *self, PyObject *args_unused)
 {
     PyEval_SetTrace((Py_tracefunc)CTracer_trace, (PyObject*)self);
-    self->started = 1;
+    self->started = TRUE;
     self->tracing_arcs = self->trace_arcs && PyObject_IsTrue(self->trace_arcs);
 
     /* start() returns a trace function usable with sys.settrace() */
@@ -957,7 +1010,7 @@ CTracer_stop(CTracer *self, PyObject *args_unused)
 {
     if (self->started) {
         PyEval_SetTrace(NULL, NULL);
-        self->started = 0;
+        self->started = FALSE;
     }
 
     Py_RETURN_NONE;
@@ -968,7 +1021,7 @@ CTracer_get_stats(CTracer *self)
 {
 #if COLLECT_STATS
     return Py_BuildValue(
-        "{sI,sI,sI,sI,sI,sI,sI,sI,si,sI,sI}",
+        "{sI,sI,sI,sI,sI,sI,sI,sI,si,sI,sI,sI}",
         "calls", self->stats.calls,
         "lines", self->stats.lines,
         "returns", self->stats.returns,
@@ -979,7 +1032,8 @@ CTracer_get_stats(CTracer *self)
         "stack_reallocs", self->stats.stack_reallocs,
         "stack_alloc", self->pdata_stack->alloc,
         "errors", self->stats.errors,
-        "pycalls", self->stats.pycalls
+        "pycalls", self->stats.pycalls,
+        "start_context_calls", self->stats.start_context_calls
         );
 #else
     Py_RETURN_NONE;
@@ -1012,6 +1066,12 @@ CTracer_members[] = {
     { "trace_arcs",         T_OBJECT, offsetof(CTracer, trace_arcs), 0,
             PyDoc_STR("Should we trace arcs, or just lines?") },
 
+    { "should_start_context", T_OBJECT, offsetof(CTracer, should_start_context), 0,
+            PyDoc_STR("Function for starting contexts.") },
+
+    { "switch_context",     T_OBJECT, offsetof(CTracer, switch_context), 0,
+            PyDoc_STR("Function for switch to a new context.") },
+
     { NULL }
 };
 
diff --git a/coverage/ctracer/tracer.h b/coverage/ctracer/tracer.h
index 053fbf62..7769b2c2 100644
--- a/coverage/ctracer/tracer.h
+++ b/coverage/ctracer/tracer.h
@@ -25,6 +25,9 @@ typedef struct CTracer {
     PyObject * file_tracers;
     PyObject * should_trace_cache;
     PyObject * trace_arcs;
+    PyObject * should_start_context;
+    PyObject * switch_context;
+    PyObject * context;
 
     /* Has the tracer been started? */
     int started;
diff --git a/coverage/ctracer/util.h b/coverage/ctracer/util.h
index 1f536194..78d9204c 100644
--- a/coverage/ctracer/util.h
+++ b/coverage/ctracer/util.h
@@ -47,4 +47,8 @@
 #define RET_OK      0
 #define RET_ERROR   -1
 
+/* Nicer booleans */
+#define FALSE   0
+#define TRUE    1
+
 #endif /* _COVERAGE_UTIL_H */