Custom writer: give full access to doc in optional Setup function (#4967)

Custom writers can specify an optional `Setup` function. The function
takes the full Pandoc document as input and should not return any value.
Users can use this function to configure the writer depending on the
given document's content or its metadata.

data/sample.lua: add sample use of Setup function.
The change allows to control the image format used to encode the image
produced from dot code.

Closes #4957

2 files changed, 50 insertions, 4 deletions

diff --git a/data/sample.lua b/data/sample.lua
index e25bf59cb..019ac13f3 100644
--- a/data/sample.lua
+++ b/data/sample.lua
@@ -2,7 +2,8 @@
 -- that is very similar to that of pandoc's HTML writer.
 -- There is one new feature: code blocks marked with class 'dot'
 -- are piped through graphviz and images are included in the HTML
--- output using 'data:' URLs.
+-- output using 'data:' URLs. The image format can be controlled
+-- via the `image_format` metadata field.
 --
 -- Invoke with: pandoc -t sample.lua
 --
@@ -13,6 +14,32 @@
 -- syntax errors.
 
 local pipe = pandoc.pipe
+local stringify = (require "pandoc.utils").stringify
+
+local image_format = "png"
+local image_mime_type = "image/png"
+
+-- Get the mime type for a given format.
+local function mime_type(img_format)
+  local formats = {
+    jpeg = "image/jpeg",
+    jpg = "image/jpeg",
+    gif = "image/gif",
+    png = "image/png",
+    svg = "image/svg+xml",
+  }
+  return formats[img_format]
+    or error("unsupported image format `" .. img_format .. "`")
+end
+
+-- Set options from document metadata.
+function Setup(doc)
+  local meta = doc.meta
+  if meta.image_format then
+    image_format = stringify(meta.image_format)
+    image_mime_type = mime_type(image_format)
+  end
+end
 
 -- Character escaping
 local function escape(s, in_attribute)
@@ -206,8 +233,8 @@ function CodeBlock(s, attr)
   -- If code block has class 'dot', pipe the contents through dot
   -- and base64, and include the base64-encoded png as a data: URL.
   if attr.class and string.match(' ' .. attr.class .. ' ',' dot ') then
-    local png = pipe("base64", {}, pipe("dot", {"-Tpng"}, s))
-    return '<img src="data:image/png;base64,' .. png .. '"/>'
+    local img = pipe("base64", {}, pipe("dot", {"-T" .. image_format}, s))
+    return '<img src="data:' .. image_mime_type .. ';base64,' .. img .. '"/>'
   -- otherwise treat as code (one could pipe through a highlighter)
   else
     return "<pre><code" .. attributes(attr) .. ">" .. escape(s) ..
@@ -325,6 +352,10 @@ end
 local meta = {}
 meta.__index =
   function(_, key)
+    -- Setup is optional, don't warn if it's not present.
+    if key == 'Setup' then
+      return
+    end
     io.stderr:write(string.format("WARNING: Undefined function '%s'\n",key))
     return function() return "" end
   end
diff --git a/src/Text/Pandoc/Writers/Custom.hs b/src/Text/Pandoc/Writers/Custom.hs
index 37fec9f0f..a5b0ed169 100644
--- a/src/Text/Pandoc/Writers/Custom.hs
+++ b/src/Text/Pandoc/Writers/Custom.hs
@@ -116,7 +116,7 @@ writeCustom luaFile opts doc@(Pandoc meta _) = do
     -- to handle this more gracefully):
     when (stat /= Lua.OK) $
       Lua.tostring' (-1) >>= throw . PandocLuaException . UTF8.toString
-    -- TODO - call hierarchicalize, so we have that info
+    runSetup doc
     rendered <- docToCustom opts doc
     context <- metaToJSON opts
                blockListToCustom
@@ -133,6 +133,21 @@ writeCustom luaFile opts doc@(Pandoc meta _) = do
               Left e  -> throw (PandocTemplateError e)
               Right r -> return (pack r)
 
+-- | Try to call a setup function. The function, if it exists, is passed the
+-- full pandoc document as parameter. This allows users to setup the writer
+-- depending on the content of the document. Accessing information on the
+-- document hierarchy is possible via the `pandoc.utils.hierarchicalize`
+-- function.
+runSetup :: Pandoc -> Lua ()
+runSetup doc = do
+  Lua.getglobal "Setup"
+  setup <- Lua.ltype Lua.stackTop
+  if setup /= Lua.TypeFunction
+    then Lua.pop 1
+    else do
+      Lua.push doc
+      Lua.call 1 0
+
 docToCustom :: WriterOptions -> Pandoc -> Lua String
 docToCustom opts (Pandoc (Meta metamap) blocks) = do
   body <- blockListToCustom blocks