1import errno
2import glob
3import os
4import shutil
5import subprocess
6import tempfile
7import time
8import uuid
9
10from docutils import nodes
11from docutils.parsers.rst import Directive
12
13
14class PythonScriptDirective(Directive):
15    """Execute an inline python script and display images.
16
17    This uses exec to execute an inline python script, copies
18    any images produced by the script, and embeds them in the document
19    along with the script.
20
21    """
22
23    required_arguments = 0
24    optional_arguments = 0
25    has_content = True
26
27    def run(self):
28        cwd = os.getcwd()
29        tmpdir = tempfile.mkdtemp()
30        os.chdir(tmpdir)
31
32        rst_file = self.state_machine.document.attributes["source"]
33        rst_dir = os.path.abspath(os.path.dirname(rst_file))
34
35        image_dir, image_rel_dir = make_image_dir(setup, rst_dir)
36
37        # Construct script from cell content
38        content = "\n".join(self.content)
39        with open("temp.py", "w") as f:
40            f.write(content)
41
42        # Use sphinx logger?
43        uid = uuid.uuid4().hex[:8]
44        print("")
45        print(f">> Contents of the script: {uid}")
46        print(content)
47        print("")
48
49        start = time.time()
50        subprocess.call(["python", "temp.py"])
51        print(f">> The execution of the script {uid} took {time.time() - start:f} s")
52        text = ""
53        for im in sorted(glob.glob("*.png")):
54            text += get_image_tag(im, image_dir, image_rel_dir)
55
56        code = content
57
58        literal = nodes.literal_block(code, code)
59        literal["language"] = "python"
60
61        attributes = {"format": "html"}
62        img_node = nodes.raw("", text, **attributes)
63
64        # clean up
65        os.chdir(cwd)
66        shutil.rmtree(tmpdir, True)
67
68        return [literal, img_node]
69
70
71def setup(app):
72    app.add_directive("python-script", PythonScriptDirective)
73    setup.app = app
74    setup.config = app.config
75    setup.confdir = app.confdir
76
77    retdict = dict(version="0.1", parallel_read_safe=True, parallel_write_safe=True)
78
79    return retdict
80
81
82def get_image_tag(filename, image_dir, image_rel_dir):
83    my_uuid = uuid.uuid4().hex
84    shutil.move(filename, image_dir + os.path.sep + my_uuid + filename)
85    relative_filename = image_rel_dir + os.path.sep + my_uuid + filename
86    return f'<img src="{relative_filename}" width="600"><br>'
87
88
89def make_image_dir(setup, rst_dir):
90    image_dir = setup.app.builder.outdir + os.path.sep + "_images"
91    rel_dir = os.path.relpath(setup.confdir, rst_dir)
92    image_rel_dir = rel_dir + os.path.sep + "_images"
93    thread_safe_mkdir(image_dir)
94    return image_dir, image_rel_dir
95
96
97def thread_safe_mkdir(dirname):
98    try:
99        os.makedirs(dirname)
100    except OSError as e:
101        if e.errno != errno.EEXIST:
102            raise
103