Crazy-ugly start to extensions for Django and Mako

--HG--
branch : django

7 files changed, 114 insertions, 44 deletions

diff --git a/coverage/codeunit.py b/coverage/codeunit.py
index 59382c23..35167a72 100644
--- a/coverage/codeunit.py
+++ b/coverage/codeunit.py
@@ -10,13 +10,15 @@ from coverage.phystokens import source_token_lines, source_encoding
 from coverage.django import DjangoTracer
 
 
-def code_unit_factory(morfs, file_locator):
+def code_unit_factory(morfs, file_locator, get_ext=None):
     """Construct a list of CodeUnits from polymorphic inputs.
 
     `morfs` is a module or a filename, or a list of same.
 
     `file_locator` is a FileLocator that can help resolve filenames.
 
+    `get_ext` TODO
+
     Returns a list of CodeUnit objects.
 
     """
@@ -28,19 +30,24 @@ def code_unit_factory(morfs, file_locator):
 
     code_units = []
     for morf in morfs:
-        # Hacked-in Mako support. Define COVERAGE_MAKO_PATH as a fragment of
-        # the path that indicates the Python file is actually a compiled Mako
-        # template. THIS IS TEMPORARY!
-        MAKO_PATH = os.environ.get('COVERAGE_MAKO_PATH')
-        if MAKO_PATH and isinstance(morf, string_class) and MAKO_PATH in morf:
-            # Super hack! Do mako both ways!
-            if 0:
-                cu = PythonCodeUnit(morf, file_locator)
-                cu.name += '_fako'
-                code_units.append(cu)
-            klass = MakoCodeUnit
-        elif isinstance(morf, string_class) and morf.endswith(".html"):
-            klass = DjangoCodeUnit
+        ext = None
+        if isinstance(morf, string_class) and get_ext:
+            ext = get_ext(morf)
+        if ext:
+            klass = DjangoTracer # NOT REALLY! TODO
+            # Hacked-in Mako support. Define COVERAGE_MAKO_PATH as a fragment of
+            # the path that indicates the Python file is actually a compiled Mako
+            # template. THIS IS TEMPORARY!
+            #MAKO_PATH = os.environ.get('COVERAGE_MAKO_PATH')
+            #if MAKO_PATH and isinstance(morf, string_class) and MAKO_PATH in morf:
+            #    # Super hack! Do mako both ways!
+            #    if 0:
+            #        cu = PythonCodeUnit(morf, file_locator)
+            #        cu.name += '_fako'
+            #        code_units.append(cu)
+            #    klass = MakoCodeUnit
+            #elif isinstance(morf, string_class) and morf.endswith(".html"):
+            #    klass = DjangoCodeUnit
         else:
             klass = PythonCodeUnit
         code_units.append(klass(morf, file_locator))
diff --git a/coverage/collector.py b/coverage/collector.py
index e3f4f630..1c530c7a 100644
--- a/coverage/collector.py
+++ b/coverage/collector.py
@@ -41,18 +41,21 @@ class PyTracer(object):
     # used to force the use of this tracer.
 
     def __init__(self):
+        # Attributes set from the collector:
         self.data = None
+        self.arcs = False
         self.should_trace = None
         self.should_trace_cache = None
         self.warn = None
-        self.handler = None
+        self.extensions = None
+
+        self.extension = None
         self.cur_file_data = None
         self.last_line = 0
         self.data_stack = []
         self.data_stacks = collections.defaultdict(list)
         self.last_exc_back = None
         self.last_exc_firstlineno = 0
-        self.arcs = False
         self.thread = None
         self.stopped = False
         self.coroutine_id_func = None
@@ -86,23 +89,24 @@ class PyTracer(object):
             if self.coroutine_id_func:
                 self.data_stack = self.data_stacks[self.coroutine_id_func()]
                 self.last_coroutine = self.coroutine_id_func()
-            self.data_stack.append((self.handler, self.cur_file_data, self.last_line))
+            self.data_stack.append((self.extension, self.cur_file_data, self.last_line))
             filename = frame.f_code.co_filename
-            if filename not in self.should_trace_cache:
+            disp = self.should_trace_cache.get(filename)
+            if disp is None:
                 disp = self.should_trace(filename, frame)
                 self.should_trace_cache[filename] = disp
-            else:
-                disp = self.should_trace_cache[filename]
             #print("called, stack is %d deep, tracename is %r" % (
             #               len(self.data_stack), tracename))
             tracename = disp.filename
-            if tracename and disp.handler:
-                tracename = disp.handler.file_name(frame)
+            if tracename and disp.extension:
+                tracename = disp.extension.file_name(frame)
             if tracename:
                 if tracename not in self.data:
                     self.data[tracename] = {}
+                    if disp.extension:
+                        self.extensions[tracename] = disp.extension.__name__
                 self.cur_file_data = self.data[tracename]
-                self.handler = disp.handler
+                self.extension = disp.extension
             else:
                 self.cur_file_data = None
             # Set the last_line to -1 because the next arc will be entering a
@@ -112,8 +116,8 @@ class PyTracer(object):
             # Record an executed line.
             #if self.coroutine_id_func:
             #    assert self.last_coroutine == self.coroutine_id_func()
-            if self.handler:
-                lineno_from, lineno_to = self.handler.line_number_range(frame)
+            if self.extension:
+                lineno_from, lineno_to = self.extension.line_number_range(frame)
             else:
                 lineno_from, lineno_to = frame.f_lineno, frame.f_lineno
             if lineno_from != -1:
@@ -134,7 +138,7 @@ class PyTracer(object):
             if self.coroutine_id_func:
                 self.data_stack = self.data_stacks[self.coroutine_id_func()]
                 self.last_coroutine = self.coroutine_id_func()
-            self.handler, self.cur_file_data, self.last_line = self.data_stack.pop()
+            self.extension, self.cur_file_data, self.last_line = self.data_stack.pop()
             #print("returned, stack is %d deep" % (len(self.data_stack)))
         elif event == 'exception':
             #print("exc", self.last_line, frame.f_lineno)
@@ -251,6 +255,8 @@ class Collector(object):
         # or mapping filenames to dicts with linenumber pairs as keys.
         self.data = {}
 
+        self.extensions = {}
+
         # A cache of the results from should_trace, the decision about whether
         # to trace execution in a file. A dict of filename to (filename or
         # None).
@@ -269,6 +275,8 @@ class Collector(object):
         tracer.warn = self.warn
         if hasattr(tracer, 'coroutine_id_func'):
             tracer.coroutine_id_func = self.coroutine_id_func
+        if hasattr(tracer, 'extensions'):
+            tracer.extensions = self.extensions
         fn = tracer.start()
         self.tracers.append(tracer)
         return fn
@@ -367,10 +375,7 @@ class Collector(object):
             # to show line data.
             line_data = {}
             for f, arcs in self.data.items():
-                line_data[f] = ldf = {}
-                for l1, _ in list(arcs.keys()):
-                    if l1:
-                        ldf[l1] = None
+                line_data[f] = dict((l1, None) for l1, _ in arcs.keys() if l1)
             return line_data
         else:
             return self.data
@@ -388,3 +393,6 @@ class Collector(object):
             return self.data
         else:
             return {}
+
+    def get_extension_data(self):
+        return self.extensions
diff --git a/coverage/config.py b/coverage/config.py
index 372439a6..e5e35856 100644
--- a/coverage/config.py
+++ b/coverage/config.py
@@ -122,6 +122,7 @@ class CoverageConfig(object):
         self.timid = False
         self.source = None
         self.debug = []
+        self.extensions = []
 
         # Defaults for [report]
         self.exclude_list = DEFAULT_EXCLUDE[:]
@@ -153,7 +154,7 @@ class CoverageConfig(object):
         if env:
             self.timid = ('--timid' in env)
 
-    MUST_BE_LIST = ["omit", "include", "debug"]
+    MUST_BE_LIST = ["omit", "include", "debug", "extensions"]
 
     def from_args(self, **kwargs):
         """Read config values from `kwargs`."""
@@ -185,12 +186,21 @@ class CoverageConfig(object):
                 self.paths[option] = cp.getlist('paths', option)
 
     CONFIG_FILE_OPTIONS = [
+        # These are *args for set_attr_from_config_option:
+        #   (attr, where, type_="")
+        #
+        #   attr is the attribute to set on the CoverageConfig object.
+        #   where is the section:name to read from the configuration file.
+        #   type_ is the optional type to apply, by using .getTYPE to read the
+        #       configuration value from the file.
+
         # [run]
         ('branch', 'run:branch', 'boolean'),
         ('coroutine', 'run:coroutine'),
         ('cover_pylib', 'run:cover_pylib', 'boolean'),
         ('data_file', 'run:data_file'),
         ('debug', 'run:debug', 'list'),
+        ('extensions', 'run:extensions', 'list'),
         ('include', 'run:include', 'list'),
         ('omit', 'run:omit', 'list'),
         ('parallel', 'run:parallel', 'boolean'),
diff --git a/coverage/control.py b/coverage/control.py
index d6bd6092..19b68ca0 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -9,6 +9,7 @@ from coverage.collector import Collector
 from coverage.config import CoverageConfig
 from coverage.data import CoverageData
 from coverage.debug import DebugControl
+from coverage.extension import load_extensions
 from coverage.files import FileLocator, TreeMatcher, FnmatchMatcher
 from coverage.files import PathAliases, find_python_files, prep_patterns
 from coverage.html import HtmlReporter
@@ -18,8 +19,6 @@ from coverage.results import Analysis, Numbers
 from coverage.summary import SummaryReporter
 from coverage.xmlreport import XmlReporter
 
-from coverage.django import DjangoTracer
-
 
 # Pypy has some unusual stuff in the "stdlib".  Consider those locations
 # when deciding where the stdlib is.
@@ -128,6 +127,10 @@ class coverage(object):
         # Create and configure the debugging controller.
         self.debug = DebugControl(self.config.debug, debug_file or sys.stderr)
 
+        # Load extensions
+        tracer_classes = load_extensions(self.config.extensions, "tracer")
+        self.tracer_extensions = [cls() for cls in tracer_classes]
+
         self.auto_data = auto_data
 
         # _exclude_re is a dict mapping exclusion list names to compiled
@@ -155,8 +158,6 @@ class coverage(object):
             coroutine=self.config.coroutine,
             )
 
-        self.django_tracer = DjangoTracer()  # should this be a class? Singleton...
-
         # Suffixes are a bit tricky.  We want to use the data suffix only when
         # collecting data, not when combining data.  So we save it as
         # `self.run_suffix` now, and promote it to `self.data_suffix` if we
@@ -274,11 +275,12 @@ class coverage(object):
 
         canonical = self.file_locator.canonical_filename(filename)
 
-        # DJANGO HACK
-        if self.django_tracer.should_trace(canonical):
-            disp.filename = canonical
-            disp.handler = self.django_tracer
-            return disp
+        # Try the extensions, see if they have an opinion about the file.
+        for tracer in self.tracer_extensions:
+            ext_disp = tracer.should_trace(canonical)
+            if ext_disp:
+                ext_disp.extension = tracer
+                return ext_disp
 
         # If the user specified source or include, then that's authoritative
         # about the outer bound of what to measure and we don't have to apply
@@ -535,8 +537,10 @@ class coverage(object):
         if not self._measured:
             return
 
+        # TODO: seems like this parallel structure is getting kinda old...
         self.data.add_line_data(self.collector.get_line_data())
         self.data.add_arc_data(self.collector.get_arc_data())
+        self.data.add_extension_data(self.collector.get_extension_data())
         self.collector.reset()
 
         # If there are still entries in the source_pkgs list, then we never
@@ -604,7 +608,8 @@ class coverage(object):
         """
         self._harvest_data()
         if not isinstance(it, CodeUnit):
-            it = code_unit_factory(it, self.file_locator)[0]
+            get_ext = self.data.extension_data().get
+            it = code_unit_factory(it, self.file_locator, get_ext)[0]
 
         return Analysis(self, it)
 
@@ -776,10 +781,10 @@ class FileDisposition(object):
         self.original_filename = original_filename
         self.filename = None
         self.reason = ""
-        self.handler = None
+        self.extension = None
 
     def nope(self, reason):
-        """A helper for return a NO answer from should_trace."""
+        """A helper for returning a NO answer from should_trace."""
         self.reason = reason
         return self
 
diff --git a/coverage/data.py b/coverage/data.py
index 042b6405..b78c931d 100644
--- a/coverage/data.py
+++ b/coverage/data.py
@@ -21,6 +21,11 @@ class CoverageData(object):
         * arcs: a dict mapping filenames to sorted lists of line number pairs:
             { 'file1': [(17,23), (17,25), (25,26)], ... }
 
+        * extensions: a dict mapping filenames to extension names:
+            { 'file1': "django.coverage", ... }
+            # TODO: how to handle the difference between a extension module
+            # name, and the class in the module?
+
     """
 
     def __init__(self, basename=None, collector=None, debug=None):
@@ -64,6 +69,14 @@ class CoverageData(object):
         #
         self.arcs = {}
 
+        # A map from canonical source file name to an extension module name:
+        #
+        #   {
+        #       'filename1.py': 'django.coverage',
+        #       ...
+        #       }
+        self.extensions = {}
+
     def usefile(self, use_file=True):
         """Set whether or not to use a disk file for data."""
         self.use_file = use_file
@@ -110,6 +123,9 @@ class CoverageData(object):
             (f, sorted(amap.keys())) for f, amap in iitems(self.arcs)
             )
 
+    def extension_data(self):
+        return self.extensions
+
     def write_file(self, filename):
         """Write the coverage data to `filename`."""
 
@@ -213,6 +229,9 @@ class CoverageData(object):
         for filename, arcs in iitems(arc_data):
             self.arcs.setdefault(filename, {}).update(arcs)
 
+    def add_extension_data(self, extension_data):
+        self.extensions.update(extension_data)
+
     def touch_file(self, filename):
         """Ensure that `filename` appears in the data, empty if needed."""
         self.lines.setdefault(filename, {})
diff --git a/coverage/extension.py b/coverage/extension.py
new file mode 100644
index 00000000..8c89b88e
--- /dev/null
+++ b/coverage/extension.py
@@ -0,0 +1,20 @@
+"""Extension management for coverage.py"""
+
+def load_extensions(modules, name):
+    """Load extensions from `modules`, finding them by `name`.
+
+    Yields the loaded extensions.
+
+    """
+
+    for module in modules:
+        try:
+            __import__(module)
+            mod = sys.modules[module]
+        except ImportError:
+            blah()
+            continue
+
+        entry = getattr(mod, name, None)
+        if entry:
+            yield entry
diff --git a/coverage/report.py b/coverage/report.py
index 03e9122c..7627d1aa 100644
--- a/coverage/report.py
+++ b/coverage/report.py
@@ -33,7 +33,8 @@ class Reporter(object):
         """
         morfs = morfs or self.coverage.data.measured_files()
         file_locator = self.coverage.file_locator
-        self.code_units = code_unit_factory(morfs, file_locator)
+        get_ext = self.coverage.data.extension_data().get
+        self.code_units = code_unit_factory(morfs, file_locator, get_ext)
 
         if self.config.include:
             patterns = prep_patterns(self.config.include)