Custom writer: provide PANDOC_DOCUMENT instead of Setup function

Custom writers have access to the global variable `PANDOC_DOCUMENT`. The
variable contains a userdata wrapper around the full pandoc AST and
exposes two fields, `meta` and `blocks`. The field content is only
marshaled on-demand, performance of scripts not accessing the fields
remains unaffected.

2 files changed, 39 insertions, 41 deletions

diff --git a/data/sample.lua b/data/sample.lua
index 019ac13f3..9d6bf0fc7 100644
--- a/data/sample.lua
+++ b/data/sample.lua
@@ -16,30 +16,24 @@
 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 = {
+-- The global variable PANDOC_DOCUMENT contains the full AST of
+-- the document which is going to be written. It can be used to
+-- configure the writer.
+local meta = PANDOC_DOCUMENT.meta
+
+-- Chose the image format based on the value of the
+-- `image_format` meta value.
+local image_format = meta.image_format
+  and stringify(meta.image_format)
+  or "png"
+local image_mime_type = ({
     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
+  })[image_format]
+  or error("unsupported image format `" .. img_format .. "`")
 
 -- Character escaping
 local function escape(s, in_attribute)
@@ -352,10 +346,6 @@ 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 a5b0ed169..3ec8781be 100644
--- a/src/Text/Pandoc/Writers/Custom.hs
+++ b/src/Text/Pandoc/Writers/Custom.hs
@@ -36,18 +36,21 @@ import Control.Arrow ((***))
 import Control.Exception
 import Control.Monad (when)
 import Data.Char (toLower)
+import Data.Data (Data)
 import Data.List (intersperse)
 import qualified Data.Map as M
 import Data.Text (Text, pack)
 import Data.Typeable
-import Foreign.Lua (Lua, Pushable)
+import Foreign.Lua (Lua, Peekable, Pushable)
+import Foreign.Lua.Userdata ( ensureUserdataMetatable, pushAnyWithMetatable
+                            , metatableName)
 import Text.Pandoc.Class (PandocIO)
 import Text.Pandoc.Definition
 import Text.Pandoc.Error
 import Text.Pandoc.Lua.Init (LuaException (LuaException), runPandocLua,
                              registerScriptPath)
 import Text.Pandoc.Lua.StackInstances ()
-import Text.Pandoc.Lua.Util (addField, dofileWithTraceback)
+import Text.Pandoc.Lua.Util (addField, addFunction, dofileWithTraceback)
 import Text.Pandoc.Options
 import Text.Pandoc.Templates
 import qualified Text.Pandoc.UTF8 as UTF8
@@ -106,17 +109,37 @@ data PandocLuaException = PandocLuaException String
 
 instance Exception PandocLuaException
 
+-- | Readonly and lazy pandoc objects.
+newtype LazyPandoc = LazyPandoc Pandoc
+  deriving (Data)
+
+instance Pushable LazyPandoc where
+  push lazyDoc = pushAnyWithMetatable pushPandocMetatable lazyDoc
+   where
+    pushPandocMetatable = ensureUserdataMetatable (metatableName lazyDoc) $
+                          addFunction "__index" indexLazyPandoc
+
+instance Peekable LazyPandoc where
+  peek = Lua.peekAny
+
+indexLazyPandoc :: LazyPandoc -> String -> Lua Lua.NumResults
+indexLazyPandoc (LazyPandoc (Pandoc meta blks)) field = 1 <$
+  case field of
+    "blocks" -> Lua.push blks
+    "meta"   -> Lua.push meta
+    _        -> Lua.pushnil
+
 -- | Convert Pandoc to custom markup.
 writeCustom :: FilePath -> WriterOptions -> Pandoc -> PandocIO Text
 writeCustom luaFile opts doc@(Pandoc meta _) = do
   res <- runPandocLua $ do
+    Lua.push (LazyPandoc doc) *> Lua.setglobal "PANDOC_DOCUMENT"
     registerScriptPath luaFile
     stat <- dofileWithTraceback luaFile
     -- check for error in lua script (later we'll change the return type
     -- to handle this more gracefully):
     when (stat /= Lua.OK) $
       Lua.tostring' (-1) >>= throw . PandocLuaException . UTF8.toString
-    runSetup doc
     rendered <- docToCustom opts doc
     context <- metaToJSON opts
                blockListToCustom
@@ -133,21 +156,6 @@ 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