doc/sphinxext/ipython_console_highlighting.py

   1 """reST directive for syntax-highlighting ipython interactive sessions.
   2 """
   3
   4 #-----------------------------------------------------------------------------
   5 # Needed modules
   6
   7 # Standard library
   8 import re
   9
  10 # Third party
  11 from pygments.lexer import Lexer, do_insertions
  12 from pygments.lexers.agile import (PythonConsoleLexer, PythonLexer,
  13                                    PythonTracebackLexer)
  14 from pygments.token import Comment, Generic
  15
  16 from sphinx import highlighting
  17
  18
  19 #-----------------------------------------------------------------------------
  20 # Global constants
  21 line_re = re.compile('.*?\n')
  22
  23 #-----------------------------------------------------------------------------
  24 # Code begins - classes and functions
  25
  26 class IPythonConsoleLexer(Lexer):
  27     """
  28     For IPython console output or doctests, such as:
  29
  30     .. sourcecode:: ipython
  31
  32       In [1]: a = 'foo'
  33
  34       In [2]: a
  35       Out[2]: 'foo'
  36
  37       In [3]: print a
  38       foo
  39
  40       In [4]: 1 / 0
  41
  42     Notes:
  43
  44       - Tracebacks are not currently supported.
  45
  46       - It assumes the default IPython prompts, not customized ones.
  47     """
  48
  49     name = 'IPython console session'
  50     aliases = ['ipython']
  51     mimetypes = ['text/x-ipython-console']
  52     input_prompt = re.compile("(In \[[0-9]+\]: )|(   \.\.\.+:)")
  53     output_prompt = re.compile("(Out\[[0-9]+\]: )|(   \.\.\.+:)")
  54     continue_prompt = re.compile("   \.\.\.+:")
  55     tb_start = re.compile("\-+")
  56
  57     def get_tokens_unprocessed(self, text):
  58         pylexer = PythonLexer(**self.options)
  59         tblexer = PythonTracebackLexer(**self.options)
  60
  61         curcode = ''
  62         insertions = []
  63         for match in line_re.finditer(text):
  64             line = match.group()
  65             input_prompt = self.input_prompt.match(line)
  66             continue_prompt = self.continue_prompt.match(line.rstrip())
  67             output_prompt = self.output_prompt.match(line)
  68             if line.startswith("#"):
  69                 insertions.append((len(curcode),
  70                                    [(0, Comment, line)]))
  71             elif input_prompt is not None:
  72                 insertions.append((len(curcode),
  73                                    [(0, Generic.Prompt, input_prompt.group())]))
  74                 curcode += line[input_prompt.end():]
  75             elif continue_prompt is not None:
  76                 insertions.append((len(curcode),
  77                                    [(0, Generic.Prompt, continue_prompt.group())]))
  78                 curcode += line[continue_prompt.end():]
  79             elif output_prompt is not None:
  80                 insertions.append((len(curcode),
  81                                    [(0, Generic.Output, output_prompt.group())]))
  82                 curcode += line[output_prompt.end():]
  83             else:
  84                 if curcode:
  85                     for item in do_insertions(insertions,
  86                                               pylexer.get_tokens_unprocessed(curcode)):
  87                         yield item
  88                         curcode = ''
  89                         insertions = []
  90                 yield match.start(), Generic.Output, line
  91         if curcode:
  92             for item in do_insertions(insertions,
  93                                       pylexer.get_tokens_unprocessed(curcode)):
  94                 yield item
  95
  96 #-----------------------------------------------------------------------------
  97 # Register the extension as a valid pygments lexer
  98 highlighting.lexers['ipython'] = IPythonConsoleLexer()