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.

1 files changed, 26 insertions, 18 deletions

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