clean code && add arguments & options

sphinx_d2lang/d2lang.py

@@ -1,27 +1,48 @@
 from docutils import nodes
 
-from docutils.parsers.rst import Directive
+from docutils.parsers.rst import Directive, directives
+from sphinx.util.docutils import SphinxDirective
 
 from shutil import which
 import shlex
 import subprocess
 import tempfile
 
-class D2langDirective(Directive):
+class D2langDirective(SphinxDirective):
+    required_arguments = 0
     has_content = True
+    optional_arguments = 3
+    option_spec = {
+        'layout': directives.unchanged_required,
+        'filename': directives.unchanged_required,
+    }
     def run(self):
-        # TODO : add tempfile for diag_source
         d2_bin = which('d2')
-        outdir = self.state.document.settings.env.app.builder.outdir
         srcdir = self.state.document.settings.env.srcdir
         diag_source = self.content
-        with tempfile.NamedTemporaryFile() as fp:
-            for line in self.content.data:
-                fp.write(bytes(line,'utf-8'))
-                fp.write(bytes('\n','utf-8'))
-            fp.seek(0)
-            cmd_line = '%s %s %s' % (d2_bin, fp.name, srcdir + "/test.svg")
-            args = shlex.split(cmd_line)
-            subprocess.run(args)
-        image_node = nodes.image(uri="/test.svg")
+        if 'filename' in self.options:
+            output_fname = self.options.get('filename')
+        else:
+            output_fname = 'test.svg'
+        if 'layout':
+            layout = self.options.get('layout')
+        else:
+            layout = 'dagre'
+        if self.arguments:
+            build_svg(self.arguments[0], srcdir, output_fname, layout)
+        else:
+            with tempfile.NamedTemporaryFile() as fp:
+                for line in self.content.data:
+                    fp.write(bytes(line,'utf-8'))
+                    fp.write(bytes('\n','utf-8'))
+                fp.seek(0)
+                build_svg(fp.name, srcdir, output_fname, layout)
+        image_node = nodes.image(uri=output_fname)
         return [image_node]
+
+def build_svg(diag_src, out_dir, filename, layout)
+    d2_bin = which('d2')
+    cmd_line = '%s -l %s %s %s' % (d2_bin, layout, diag_src, out_dir + "/" + filename)
+    args = shlex.split(cmd_line)
+    subprocess.run(args)
+    return True