Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
yt-project
GitHub Repository: yt-project/yt
 Path: blob/main/doc/extensions/yt_cookbook.py
928 views
1
# This extension is quite simple:

2
#  1. It accepts a script name

3
#  2. This script is added to the document in a literalinclude

4
#  3. Any _static images found will be added

5


6
import glob

7
import os

8
import shutil

9


10
from docutils.parsers.rst import Directive, directives

11


12
# Some of this magic comes from the matplotlib plot_directive.

13


14


15
def setup(app):

16
    app.add_directive("yt_cookbook", CookbookScript)

17
    setup.app = app

18
    setup.config = app.config

19
    setup.confdir = app.confdir

20


21
    retdict = dict(version="0.1", parallel_read_safe=True, parallel_write_safe=True)

22


23
    return retdict

24


25


26
data_patterns = ["*.h5", "*.out", "*.dat", "*.mp4"]

27


28


29
class CookbookScript(Directive):

30
    required_arguments = 1

31
    optional_arguments = 0

32


33
    def run(self):

34
        rst_file = self.state_machine.document.attributes["source"]

35
        rst_dir = os.path.abspath(os.path.dirname(rst_file))

36
        script_fn = directives.path(self.arguments[0])

37
        script_bn = os.path.basename(script_fn)

38
        script_name = os.path.basename(self.arguments[0]).split(".")[0]

39


40
        # This magic is from matplotlib

41
        dest_dir = os.path.abspath(

42
            os.path.join(setup.app.builder.outdir, os.path.dirname(script_fn))

43
        )

44
        if not os.path.exists(dest_dir):

45
            os.makedirs(dest_dir)  # no problem here for me, but just use built-ins

46


47
        rel_dir = os.path.relpath(rst_dir, setup.confdir)

48
        place = os.path.join(dest_dir, rel_dir)

49
        if not os.path.isdir(place):

50
            os.makedirs(place)

51
        shutil.copyfile(

52
            os.path.join(rst_dir, script_fn), os.path.join(place, script_bn)

53
        )

54


55
        im_path = os.path.join(rst_dir, "_static")

56
        images = sorted(glob.glob(os.path.join(im_path, f"{script_name}__*.png")))

57
        lines = []

58
        lines.append(f"(`{script_bn} <{script_fn}>`__)")

59
        lines.append("\n")

60
        lines.append("\n")

61
        lines.append(f".. literalinclude:: {self.arguments[0]}")

62
        lines.append("\n")

63
        lines.append("\n")

64
        for im in images:

65
            im_name = os.path.join("_static", os.path.basename(im))

66
            lines.append(f".. image:: {im_name}")

67
            lines.append("   :width: 400")

68
            lines.append(f"   :target: ../_images/{os.path.basename(im)}")

69
            lines.append("\n")

70
        lines.append("\n")

71
        for ext in data_patterns:

72
            data_files = sorted(

73
                glob.glob(os.path.join(im_path, f"{script_name}__*.{ext}"))

74
            )

75
            for df in data_files:

76
                df_bn = os.path.basename(df)

77
                shutil.copyfile(

78
                    os.path.join(rst_dir, df), os.path.join(dest_dir, rel_dir, df_bn)

79
                )

80
                lines.append(f" * Data: `{df_bn} <{df}>`__)")

81
            lines.append("\n")

82
        self.state_machine.insert_input(lines, rst_file)

83
        return []

84


85