Implement ALLOW_BYTES doctest option

Fix pytest-dev#1287

_pytest/doctest.py

@@ -79,7 +79,7 @@ def repr_failure(self, excinfo):
                 lineno = test.lineno + example.lineno + 1
             message = excinfo.type.__name__
             reprlocation = ReprFileLocation(filename, lineno, message)
-            checker = _get_unicode_checker()
+            checker = _get_checker()
             REPORT_UDIFF = doctest.REPORT_UDIFF
             filelines = py.path.local(filename).readlines(cr=0)
             lines = []
@@ -118,7 +118,9 @@ def _get_flag_lookup():
                 ELLIPSIS=doctest.ELLIPSIS,
                 IGNORE_EXCEPTION_DETAIL=doctest.IGNORE_EXCEPTION_DETAIL,
                 COMPARISON_FLAGS=doctest.COMPARISON_FLAGS,
-                ALLOW_UNICODE=_get_allow_unicode_flag())
+                ALLOW_UNICODE=_get_allow_unicode_flag(),
+                ALLOW_BYTES=_get_allow_bytes_flag(),
+                )
 
 
 def get_optionflags(parent):
@@ -147,7 +149,7 @@ def runtest(self):
 
         optionflags = get_optionflags(self)
         runner = doctest.DebugRunner(verbose=0, optionflags=optionflags,
-                                     checker=_get_unicode_checker())
+                                     checker=_get_checker())
 
         parser = doctest.DocTestParser()
         test = parser.get_doctest(text, globs, name, filename, 0)
@@ -182,7 +184,7 @@ def collect(self):
         finder = doctest.DocTestFinder()
         optionflags = get_optionflags(self)
         runner = doctest.DebugRunner(verbose=0, optionflags=optionflags,
-                                     checker=_get_unicode_checker())
+                                     checker=_get_checker())
         for test in finder.find(module, module.__name__):
             if test.examples:  # skip empty doctests
                 yield DoctestItem(test.name, self, runner, test)
@@ -204,52 +206,60 @@ def func():
     return fixture_request
 
 
-def _get_unicode_checker():
+def _get_checker():
     """
     Returns a doctest.OutputChecker subclass that takes in account the
-    ALLOW_UNICODE option to ignore u'' prefixes in strings. Useful
-    when the same doctest should run in Python 2 and Python 3.
+    ALLOW_UNICODE option to ignore u'' prefixes in strings and ALLOW_BYTES
+    to strip b'' prefixes.
+    Useful when the same doctest should run in Python 2 and Python 3.
 
     An inner class is used to avoid importing "doctest" at the module
     level.
     """
-    if hasattr(_get_unicode_checker, 'UnicodeOutputChecker'):
-        return _get_unicode_checker.UnicodeOutputChecker()
+    if hasattr(_get_checker, 'LiteralsOutputChecker'):
+        return _get_checker.LiteralsOutputChecker()
 
     import doctest
     import re
 
-    class UnicodeOutputChecker(doctest.OutputChecker):
+    class LiteralsOutputChecker(doctest.OutputChecker):
         """
         Copied from doctest_nose_plugin.py from the nltk project:
             https://github.com/nltk/nltk
+
+        Further extended to also support byte literals.
         """
 
-        _literal_re = re.compile(r"(\W|^)[uU]([rR]?[\'\"])", re.UNICODE)
+        _unicode_literal_re = re.compile(r"(\W|^)[uU]([rR]?[\'\"])", re.UNICODE)
+        _bytes_literal_re = re.compile(r"(\W|^)[bB]([rR]?[\'\"])", re.UNICODE)
 
         def check_output(self, want, got, optionflags):
             res = doctest.OutputChecker.check_output(self, want, got,
                                                      optionflags)
             if res:
                 return True
 
-            if not (optionflags & _get_allow_unicode_flag()):
+            allow_unicode = optionflags & _get_allow_unicode_flag()
+            allow_bytes = optionflags & _get_allow_bytes_flag()
+            if not allow_unicode and not allow_bytes:
                 return False
 
             else:  # pragma: no cover
-                # the code below will end up executed only in Python 2 in
-                # our tests, and our coverage check runs in Python 3 only
-                def remove_u_prefixes(txt):
-                    return re.sub(self._literal_re, r'\1\2', txt)
-
-                want = remove_u_prefixes(want)
-                got = remove_u_prefixes(got)
+                def remove_prefixes(regex, txt):
+                    return re.sub(regex, r'\1\2', txt)
+
+                if allow_unicode:
+                    want = remove_prefixes(self._unicode_literal_re, want)
+                    got = remove_prefixes(self._unicode_literal_re, got)
+                if allow_bytes:
+                    want = remove_prefixes(self._bytes_literal_re, want)
+                    got = remove_prefixes(self._bytes_literal_re, got)
                 res = doctest.OutputChecker.check_output(self, want, got,
                                                          optionflags)
                 return res
 
-    _get_unicode_checker.UnicodeOutputChecker = UnicodeOutputChecker
-    return _get_unicode_checker.UnicodeOutputChecker()
+    _get_checker.LiteralsOutputChecker = LiteralsOutputChecker
+    return _get_checker.LiteralsOutputChecker()
 
 
 def _get_allow_unicode_flag():
@@ -258,3 +268,11 @@ def _get_allow_unicode_flag():
     """
     import doctest
     return doctest.register_optionflag('ALLOW_UNICODE')
+
+
+def _get_allow_bytes_flag():
+    """
+    Registers and returns the ALLOW_BYTES flag.
+    """
+    import doctest
+    return doctest.register_optionflag('ALLOW_BYTES')

testing/test_doctest.py

@@ -371,6 +371,9 @@ def foo():
                                     "--junit-xml=junit.xml")
         reprec.assertoutcome(failed=1)
 
+
+class TestLiterals:
+
     @pytest.mark.parametrize('config_mode', ['ini', 'comment'])
     def test_allow_unicode(self, testdir, config_mode):
         """Test that doctests which output unicode work in all python versions
@@ -400,6 +403,35 @@ def foo():
         reprec = testdir.inline_run("--doctest-modules")
         reprec.assertoutcome(passed=2)
 
+    @pytest.mark.parametrize('config_mode', ['ini', 'comment'])
+    def test_allow_bytes(self, testdir, config_mode):
+        """Test that doctests which output bytes work in all python versions
+        tested by pytest when the ALLOW_BYTES option is used (either in
+        the ini file or by an inline comment)(#1287).
+        """
+        if config_mode == 'ini':
+            testdir.makeini('''
+            [pytest]
+            doctest_optionflags = ALLOW_BYTES
+            ''')
+            comment = ''
+        else:
+            comment = '#doctest: +ALLOW_BYTES'
+
+        testdir.maketxtfile(test_doc="""
+            >>> b'foo'  {comment}
+            'foo'
+        """.format(comment=comment))
+        testdir.makepyfile(foo="""
+            def foo():
+              '''
+              >>> b'foo'  {comment}
+              'foo'
+              '''
+        """.format(comment=comment))
+        reprec = testdir.inline_run("--doctest-modules")
+        reprec.assertoutcome(passed=2)
+
     def test_unicode_string(self, testdir):
         """Test that doctests which output unicode fail in Python 2 when
         the ALLOW_UNICODE option is not used. The same test should pass
@@ -413,6 +445,19 @@ def test_unicode_string(self, testdir):
         passed = int(sys.version_info[0] >= 3)
         reprec.assertoutcome(passed=passed, failed=int(not passed))
 
+    def test_bytes_literal(self, testdir):
+        """Test that doctests which output bytes fail in Python 3 when
+        the ALLOW_BYTES option is not used. The same test should pass
+        in Python 2 (#1287).
+        """
+        testdir.maketxtfile(test_doc="""
+            >>> b'foo'
+            'foo'
+        """)
+        reprec = testdir.inline_run()
+        passed = int(sys.version_info[0] == 2)
+        reprec.assertoutcome(passed=passed, failed=int(not passed))
+
 
 class TestDoctestSkips:
     """
@@ -579,4 +624,4 @@ def auto(request):
         """)
         result = testdir.runpytest('--doctest-modules')
         assert 'FAILURES' not in str(result.stdout.str())
-        result.stdout.fnmatch_lines(['*=== 1 passed in *'])
+        result.stdout.fnmatch_lines(['*=== 1 passed in *'])