Support doctest highlight in Python lexer (#1932)

lib/rouge/lexers/python.rb

@@ -80,6 +80,8 @@ def current_string
           groups Punctuation, Text, Str::Doc
         end
 
+        rule %r/\.\.\.\B$/, Name::Builtin::Pseudo
+
         rule %r/[^\S\n]+/, Text
         rule %r(#(.*)?\n?), Comment::Single
         rule %r/[\[\]{}:(),;.]/, Punctuation
@@ -88,6 +90,8 @@ def current_string
 
         rule %r/@#{dotted_identifier}/i, Name::Decorator
 
+        rule %r/(>>>|\.\.\.)\B/, Generic::Prompt
+
         rule %r/(in|is|and|or|not)\b/, Operator::Word
         rule %r/(<<|>>|\/\/|\*\*)=?/, Operator
         rule %r/[-~+\/*%=<>&^|@]=?|!=/, Operator
@@ -120,7 +124,7 @@ def current_string
         # TODO: not in python 3
         rule %r/`.*?`/, Str::Backtick
         rule %r/([rfbu]{0,2})('''|"""|['"])/i do |m|
-          groups Str::Affix, Str
+          groups Str::Affix, Str::Heredoc
           current_string.register type: m[1].downcase, delim: m[2]
           push :generic_string
         end
@@ -178,11 +182,12 @@ def current_string
       end
 
       state :generic_string do
-        rule %r/[^'"\\{]+/, Str
+        rule %r/>>>|\.\.\./, Generic::Prompt, :doctest
+        rule %r/[^'"\\{]+?/, Str
         rule %r/{{/, Str
 
         rule %r/'''|"""|['"]/ do |m|
-          token Str
+          token Str::Heredoc
           if current_string.delim? m[0]
             current_string.remove
             pop!
@@ -220,6 +225,17 @@ def current_string
         rule %r/\\./, Str, :pop!
       end
 
+      state :doctest do
+        rule %r/\n\n/, Text, :pop!
+
+        rule %r/'''|"""/ do
+          token Str::Heredoc
+          pop!(2) if in_state?(:generic_string) # pop :doctest and :generic_string
+        end
+
+        mixin :root
+      end
+
       state :generic_interpol do
         rule %r/[^{}!:]+/ do |m|
           recurse m[0]

spec/visual/samples/python

@@ -173,3 +173,37 @@ class Spam:
     pass
 
 spam = Spam()
+
+# Doctest
+def factorial(n):
+    """Return the factorial of n, an exact integer >= 0.
+
+    >>> [factorial(n) for n in range(6)]
+    [1, 1, 2, 6, 24, 120]
+
+    >>> for i in [factorial(n) for n in range(2)]:
+    ...     print(i)
+    1
+    1
+    >>> factorial(30)
+    265252859812191058636308480000000
+    >>> factorial(-1)
+    Traceback (most recent call last):
+        ...
+    ValueError: n must be >= 0
+
+    Factorials of floats are OK, but the float must be an exact integer:
+    >>> factorial(30.1)
+    Traceback (most recent call last):
+        ...
+    ValueError: n must be exact integer
+    >>> factorial(30.0)
+    265252859812191058636308480000000
+
+    It must also not be ridiculously large:
+    >>> factorial(1e100)
+    Traceback (most recent call last):
+        ...
+    OverflowError: n too large
+    """
+    print("hello world")