gistfile1.py

from __future__ import division
import os
import sys
import codecs
import textwrap
from StringIO import StringIO
from docutils import nodes
from docutils.parsers.rst import Directive, directives
from sphinx.directives.code import LiteralInclude
from sphinx.util.nodes import set_source_info


class CodeExample(Directive):
    """
    This is like literalinclude, but add the code and the result of
    its invocation in a python prompt.
    """

    has_content = False
    required_arguments = 1
    optional_arguments = 0
    final_argument_whitespace = False
    option_spec = {
        'linenos': directives.flag,
        'tab-width': int,
        'language': directives.unchanged_required,
        'lines': directives.unchanged_required,
        'start-after': directives.unchanged_required,
        'end-before': directives.unchanged_required,
        'prepend': directives.unchanged_required,
        'append': directives.unchanged_required,
        'emphasize-lines': directives.unchanged_required,
    }

    @staticmethod
    def generate_example(filename, env):
        for path in env.config['code_add_python_path']:
            sys.path.append(path)

        prompt = ">>>"
        oldstdout = sys.stdout
        buffer = StringIO()
        sys.stdout = buffer

        global_env = {}
        local_env = {}

        with open(filename) as data:
            for line in data:
                code = line.strip('\n')
                if code.startswith("from") or code.startswith("import"):
                    eval(compile(code, '', 'exec'), global_env, local_env)
                elif code:
                    print prompt, code

                    # this is very hacky, we need to get assignments
                    # but not keyword arguments
                    if " = " in code:
                        # if an expression is an assignment we need to
                        # print the result the assignment. To do that
                        # we print the variable value. Hacky and ugly.
                        fcode = code
                        ccode = compile(fcode, filename, 'exec', division.compiler_flag)
                        eval(ccode, global_env, local_env)
                        var = code.split("=")[0]
                        print prompt, eval(var, global_env, local_env)
                    else:
                        fcode = "print({0})".format(code)
                        ccode = compile(fcode, filename, 'exec', division.compiler_flag)
                        print ">>>",
                        eval(ccode, global_env, local_env)

        sys.stdout = sys.__stdout__
        return buffer.getvalue()

    def run(self):
        document = self.state.document
        if not document.settings.file_insertion_enabled:
            return [document.reporter.warning('File insertion disabled',
                                              line=self.lineno)]
        env = document.settings.env
        full_filename = os.path.join(env.config.code_example_dir, self.arguments[0])
        rel_filename, filename = env.relfn2path(full_filename)

        if 'pyobject' in self.options and 'lines' in self.options:
            return [document.reporter.warning(
                'Cannot use both "pyobject" and "lines" options',
                line=self.lineno)]

        lines = self.generate_example(filename, env).split('\n')

        linespec = self.options.get('lines')
        if linespec is not None:
            try:
                linelist = parselinenos(linespec, len(lines))
            except ValueError, err:
                return [document.reporter.warning(str(err), line=self.lineno)]
            # just ignore nonexisting lines
            nlines = len(lines)
            lines = [lines[i] for i in linelist if i < nlines]
            if not lines:
                return [document.reporter.warning(
                    'Line spec %r: no lines pulled from include file %r' %
                    (linespec, filename), line=self.lineno)]

        linespec = self.options.get('emphasize-lines')
        if linespec:
            try:
                hl_lines = [x+1 for x in parselinenos(linespec, len(lines))]
            except ValueError, err:
                return [document.reporter.warning(str(err), line=self.lineno)]
        else:
            hl_lines = None

        startafter = self.options.get('start-after')
        endbefore  = self.options.get('end-before')
        prepend    = self.options.get('prepend')
        append     = self.options.get('append')
        if startafter is not None or endbefore is not None:
            use = not startafter
            res = []
            for line in lines:
                if not use and startafter and startafter in line:
                    use = True
                elif use and endbefore and endbefore in line:
                    use = False
                    break
                elif use:
                    res.append(line)
            lines = res

        if prepend:
           lines.insert(0, prepend + '\n')
        if append:
           lines.append(append + '\n')

        wrap_value = env.config['code_example_wrap']
        txt = []
        if wrap_value:
            for line in lines:
                if len(line) >= 75:
                    r = textwrap.fill(line, subsequent_indent='>>> ', width=wrap_value)
                else:
                    r = line
                txt.append(r)

        lines = txt if txt else lines
        text = '\n'.join(lines)
            
        if self.options.get('tab-width'):
            text = text.expandtabs(self.options['tab-width'])

        retnode = nodes.literal_block(text, text, source=filename)
        set_source_info(self, retnode)
        if self.options.get('language', ''):
            retnode['language'] = self.options['language']
        if 'linenos' in self.options:
            retnode['linenos'] = True
        if hl_lines is not None:
            retnode['highlight_args'] = {'hl_lines': hl_lines}
        env.note_dependency(rel_filename)
        return [retnode]



def setup(app):
    app.add_directive('code-example', CodeExample)
    app.add_config_value('code_example_dir', 'code-example', False)
    app.add_config_value('code_add_python_path', [], False)
    app.add_config_value('code_example_wrap', False, False)