Pragmas on decorators apply to the entire function or class. #131

4 files changed, 89 insertions, 13 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index 7ff7334d..65424168 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,16 @@ Change history for Coverage.py
 ==============================
 
 
+Version 4.1, in progress
+------------------------
+
+- Pragmas to disable coverage measurement can now be used on decorator lines,
+  and they will apply to the entire function or class being decorated.  This
+  implements the feature requested in `issue 131`_.
+
+.. _issue 131: https://bitbucket.org/ned/coveragepy/issues/131/pragma-on-a-decorator-line-should-affect
+
+
 Version 4.0.3, 24 November 2015
 -------------------------------
 
diff --git a/coverage/parser.py b/coverage/parser.py
index 111826da..7b8a60f1 100644
--- a/coverage/parser.py
+++ b/coverage/parser.py
@@ -117,9 +117,11 @@ class PythonParser(object):
         indent = 0
         exclude_indent = 0
         excluding = False
+        excluding_decorators = False
         prev_toktype = token.INDENT
         first_line = None
         empty = True
+        first_on_line = True
 
         tokgen = generate_tokens(self.text)
         for toktype, ttext, (slineno, _), (elineno, _), ltext in tokgen:
@@ -132,18 +134,29 @@ class PythonParser(object):
                 indent += 1
             elif toktype == token.DEDENT:
                 indent -= 1
-            elif toktype == token.NAME and ttext == 'class':
-                # Class definitions look like branches in the byte code, so
-                # we need to exclude them.  The simplest way is to note the
-                # lines with the 'class' keyword.
-                self.raw_classdefs.add(slineno)
-            elif toktype == token.OP and ttext == ':':
-                if not excluding and elineno in self.raw_excluded:
-                    # Start excluding a suite.  We trigger off of the colon
-                    # token so that the #pragma comment will be recognized on
-                    # the same line as the colon.
-                    exclude_indent = indent
-                    excluding = True
+            elif toktype == token.NAME:
+                if ttext == 'class':
+                    # Class definitions look like branches in the byte code, so
+                    # we need to exclude them.  The simplest way is to note the
+                    # lines with the 'class' keyword.
+                    self.raw_classdefs.add(slineno)
+            elif toktype == token.OP:
+                if ttext == ':':
+                    should_exclude = (elineno in self.raw_excluded) or excluding_decorators
+                    if not excluding and should_exclude:
+                        # Start excluding a suite.  We trigger off of the colon
+                        # token so that the #pragma comment will be recognized on
+                        # the same line as the colon.
+                        self.raw_excluded.add(elineno)
+                        exclude_indent = indent
+                        excluding = True
+                        excluding_decorators = False
+                elif ttext == '@' and first_on_line:
+                    # A decorator.
+                    if elineno in self.raw_excluded:
+                        excluding_decorators = True
+                    if excluding_decorators:
+                        self.raw_excluded.add(elineno)
             elif toktype == token.STRING and prev_toktype == token.INDENT:
                 # Strings that are first on an indented line are docstrings.
                 # (a trick from trace.py in the stdlib.) This works for
@@ -158,6 +171,7 @@ class PythonParser(object):
                     for l in range(first_line, elineno+1):
                         self._multiline[l] = first_line
                 first_line = None
+                first_on_line = True
 
             if ttext.strip() and toktype != tokenize.COMMENT:
                 # A non-whitespace token.
@@ -171,6 +185,7 @@ class PythonParser(object):
                         excluding = False
                     if excluding:
                         self.raw_excluded.add(elineno)
+                    first_on_line = False
 
             prev_toktype = toktype
 
diff --git a/lab/parser.py b/lab/parser.py
index bb593f8f..70c2b6b9 100644
--- a/lab/parser.py
+++ b/lab/parser.py
@@ -108,7 +108,9 @@ class ParserMain(object):
 
                 for lineno, ltext in enumerate(cp.lines, start=1):
                     m0 = m1 = m2 = m3 = a = ' '
-                    if lineno in cp.raw_statements:
+                    if lineno in cp.statements:
+                        m0 = '='
+                    elif lineno in cp.raw_statements:
                         m0 = '-'
                     exits = exit_counts.get(lineno, 0)
                     if exits > 1:
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 372bf79b..44a261d9 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -136,6 +136,55 @@ class PythonParserTest(CoverageTest):
                 '''
                 """)
 
+    def test_decorator_pragmas(self):
+        parser = self.parse_source("""\
+            # 1
+
+            @foo(3)                     # nocover
+            @bar
+            def func(x, y=5):
+                return 6
+
+            class Foo:      # the only statement...
+                '''9'''
+                @foo                    # nocover
+                def __init__(self):
+                    '''12'''
+                    return 13
+
+                @foo(                   # nocover
+                    16,
+                    17,
+                )
+                def meth(self):
+                    return 20
+
+            @foo(                       # nocover
+                23
+            )
+            def func(x=25):
+                return 26
+            """)
+        self.assertEqual(
+            parser.raw_statements,
+            set([3, 4, 5, 6, 8, 9, 10, 13, 15, 16, 17, 20, 22, 23, 25, 26])
+        )
+        self.assertEqual(parser.statements, set([8]))
+
+    def test_class_decorator_pragmas(self):
+        parser = self.parse_source("""\
+            class Foo(object):
+                def __init__(self):
+                    self.x = 3
+
+            @foo                        # nocover
+            class Bar(object):
+                def __init__(self):
+                    self.x = 8
+            """)
+        self.assertEqual(parser.raw_statements, set([1, 2, 3, 5, 6, 7, 8]))
+        self.assertEqual(parser.statements, set([1, 2, 3]))
+
 
 class ParserFileTest(CoverageTest):
     """Tests for coverage.py's code parsing from files."""