Merge pull request #3945 from tarleb/lua-mediabag

Lua mediabag module

6 files changed, 252 insertions, 41 deletions

diff --git a/doc/lua-filters.md b/doc/lua-filters.md
index b591a747f..d17c44adf 100644
--- a/doc/lua-filters.md
+++ b/doc/lua-filters.md
@@ -1062,3 +1062,80 @@ Lua functions for pandoc scripts.
         return {pandoc.global_filter()}
         -- the above is equivallent to
         -- return {{Str = Str}}
+
+# Submodule mediabag
+
+The submodule `mediabag` allows accessing pandoc's media
+storage. The "media bag" is used when pandoc is called with the
+`--extract-media` or `--standalone`/`-s` option.
+
+[`insert (filepath, mime_type, contents)`]{#mediabag-insert}
+
+:   Adds a new entry to pandoc's media bag.
+
+    Parameters:
+
+    `filepath`:
+    :   filename and path relative to the output folder.
+
+    `mime_type`:
+    :   the file's MIME type
+
+    `contents`:
+    :   the binary contents of the file.
+
+    Usage:
+
+        local fp = "media/hello.txt"
+        local mt = "text/plain"
+        local contents = "Hello, World!"
+        pandoc.mediabag(fp, mt, contents)
+
+[`list ()`]{#mediabag-list}
+
+:   Get a summary of the current media bag contents.
+
+    Returns: A list of elements summarizing each entry in the
+    media bag. The summary item contains the keys `path`,
+    `type`, and `length`, giving the filepath, MIME type, and
+    length of contents in bytes, respectively.
+
+    Usage:
+
+        -- calculate the size of the media bag.
+        local mb_items = pandoc.mediabag.list()
+        local sum = 0
+        for i = 1, #mb_items:
+            sum = sum + mb_items[i].length
+        end
+        print(sum)
+
+[`lookup (filepath)`]{#mediabag-lookup}
+
+:   Lookup a media item in the media bag, returning mime type
+    and contents.
+
+    Parameters:
+
+    `filepath`:
+    :   name of the file to look up.
+
+    Returns:
+
+    -   the entries MIME type, or nil if the file was not found.
+    -   contents of the file, or nil if the file was not found.
+
+    Usage:
+
+        local filename = "media/diagram.png"
+        local mt, contents = pandoc.mediabag.lookup(filename)
+
+[`fetch (source, base_url)`]{#mediabag-fetch}
+
+:   Fetches the given source and inserts it into the media bag
+    using a SHA1 hash of the content as filename.
+
+    Usage:
+
+        local diagram_url = "https://pandoc.org/diagram.jpg"
+        pandoc.mediabag.fetch(diagram_url, ".")
diff --git a/src/Text/Pandoc/App.hs b/src/Text/Pandoc/App.hs
index 206c47b30..9b3055b35 100644
--- a/src/Text/Pandoc/App.hs
+++ b/src/Text/Pandoc/App.hs
@@ -511,10 +511,10 @@ convertWithOpts opts = do
               (   (if isJust (optExtractMedia opts)
                       then fillMediaBag (writerSourceURL writerOptions)
                       else return)
-              >=> maybe return extractMedia (optExtractMedia opts)
               >=> return . flip (foldr addMetadata) metadata
-              >=> applyTransforms transforms
               >=> applyLuaFilters datadir (optLuaFilters opts) format
+              >=> maybe return extractMedia (optExtractMedia opts)
+              >=> applyTransforms transforms
               >=> applyFilters readerOpts datadir filters' [format]
               )
     media <- getMediaBag
@@ -850,16 +850,15 @@ expandFilterPath mbDatadir fp = liftIO $ do
                     else return fp
                _ -> return fp
 
-applyLuaFilters :: MonadIO m
-                => Maybe FilePath -> [FilePath] -> String -> Pandoc
-                -> m Pandoc
+applyLuaFilters :: Maybe FilePath -> [FilePath] -> String -> Pandoc
+                -> PandocIO Pandoc
 applyLuaFilters mbDatadir filters format d = do
   expandedFilters <- mapM (expandFilterPath mbDatadir) filters
-  let go f d' = liftIO $ do
-        res <- E.try (runLuaFilter mbDatadir f format d')
+  let go f d' = do
+        res <- runLuaFilter mbDatadir f format d'
         case res of
-             Right x -> return x
-             Left (LuaException s) -> E.throw (PandocFilterError f s)
+          Right x -> return x
+          Left (LuaException s) -> E.throw (PandocFilterError f s)
   foldrM ($) d $ map go expandedFilters
 
 applyFilters :: MonadIO m
diff --git a/src/Text/Pandoc/Class.hs b/src/Text/Pandoc/Class.hs
index 98c567afc..f60062d6c 100644
--- a/src/Text/Pandoc/Class.hs
+++ b/src/Text/Pandoc/Class.hs
@@ -79,6 +79,7 @@ module Text.Pandoc.Class ( PandocMonad(..)
                          , runPure
                          , readDefaultDataFile
                          , readDataFile
+                         , fetchMediaResource
                          , fillMediaBag
                          , extractMedia
                          , toLang
@@ -246,9 +247,9 @@ getMediaBag = getsCommonState stMediaBag
 
 insertMedia :: PandocMonad m => FilePath -> Maybe MimeType -> BL.ByteString -> m ()
 insertMedia fp mime bs = do
-  mb <- getsCommonState stMediaBag
+  mb <- getMediaBag
   let mb' = MB.insertMedia fp mime bs mb
-  modifyCommonState $ \st -> st{stMediaBag = mb' }
+  setMediaBag mb'
 
 getInputFiles :: PandocMonad m => m (Maybe [FilePath])
 getInputFiles = getsCommonState stInputFiles
@@ -633,6 +634,20 @@ withPaths (p:ps) action fp =
   catchError (action (p </> fp))
              (\_ -> withPaths ps action fp)
 
+-- | Fetch local or remote resource (like an image) and provide data suitable
+-- for adding it to the MediaBag.
+fetchMediaResource :: PandocMonad m
+              => Maybe String -> String
+              -> m (FilePath, Maybe MimeType, BL.ByteString)
+fetchMediaResource sourceUrl src = do
+  (bs, mt) <- downloadOrRead sourceUrl src
+  let ext = fromMaybe (takeExtension src)
+                      (mt >>= extensionFromMimeType)
+  let bs' = BL.fromChunks [bs]
+  let basename = showDigest $ sha1 bs'
+  let fname = basename <.> ext
+  return (fname, mt, bs')
+
 -- | Traverse tree, filling media bag for any images that
 -- aren't already in the media bag.
 fillMediaBag :: PandocMonad m => Maybe String -> Pandoc -> m Pandoc
@@ -643,13 +658,8 @@ fillMediaBag sourceURL d = walkM handleImage d
               case lookupMedia src mediabag of
                 Just (_, _) -> return $ Image attr lab (src, tit)
                 Nothing -> do
-                  (bs, mt) <- downloadOrRead sourceURL src
-                  let ext = fromMaybe (takeExtension src)
-                              (mt >>= extensionFromMimeType)
-                  let bs' = BL.fromChunks [bs]
-                  let basename = showDigest $ sha1 bs'
-                  let fname = basename <.> ext
-                  insertMedia fname mt bs'
+                  (fname, mt, bs) <- fetchMediaResource sourceURL src
+                  insertMedia fname mt bs
                   return $ Image attr lab (fname, tit))
           (\e ->
               case e of
diff --git a/src/Text/Pandoc/Lua.hs b/src/Text/Pandoc/Lua.hs
index ab3b5f4ca..f7e74d0a8 100644
--- a/src/Text/Pandoc/Lua.hs
+++ b/src/Text/Pandoc/Lua.hs
@@ -39,26 +39,40 @@ import Control.Monad.Trans (MonadIO (..))
 import Data.Data (DataType, Data, toConstr, showConstr, dataTypeOf,
                   dataTypeConstrs, dataTypeName, tyconUQname)
 import Data.Foldable (foldrM)
+import Data.IORef (IORef, newIORef, readIORef)
 import Data.Map (Map)
 import Data.Maybe (isJust)
 import Foreign.Lua (Lua, FromLuaStack (peek), LuaException (..), StackIndex,
                     Status (OK), ToLuaStack (push))
+import Text.Pandoc.Class (PandocIO, getMediaBag, setMediaBag)
+import Text.Pandoc.MediaBag (MediaBag)
 import Text.Pandoc.Definition
-import Text.Pandoc.Lua.PandocModule (pushPandocModule)
+import Text.Pandoc.Lua.PandocModule (pushPandocModule, pushMediaBagModule)
 import Text.Pandoc.Walk (walkM)
 
 import qualified Data.Map as Map
 import qualified Foreign.Lua as Lua
 
-runLuaFilter :: (MonadIO m)
-             => Maybe FilePath -> FilePath -> String -> Pandoc -> m Pandoc
-runLuaFilter datadir filterPath format pd = liftIO . Lua.runLua $ do
+runLuaFilter :: Maybe FilePath -> FilePath -> String
+             -> Pandoc -> PandocIO (Either LuaException Pandoc)
+runLuaFilter datadir filterPath format pd = do
+  mediaBag <- getMediaBag
+  mediaBagRef <- liftIO (newIORef mediaBag)
+  res <- liftIO . Lua.runLuaEither $
+         runLuaFilter' datadir filterPath format mediaBagRef pd
+  newMediaBag <- liftIO (readIORef mediaBagRef)
+  setMediaBag newMediaBag
+  return res
+
+runLuaFilter' :: Maybe FilePath -> FilePath -> String -> IORef MediaBag
+              -> Pandoc -> Lua Pandoc
+runLuaFilter' datadir filterPath format mbRef pd = do
   Lua.openlibs
   -- store module in global "pandoc"
   pushPandocModule datadir
   Lua.setglobal "pandoc"
-  push format
-  Lua.setglobal "FORMAT"
+  addMediaBagModule
+  registerFormat
   top <- Lua.gettop
   stat <- Lua.dofile filterPath
   if stat /= OK
@@ -71,6 +85,16 @@ runLuaFilter datadir filterPath format pd = liftIO . Lua.runLua $ do
       when (newtop - top < 1) pushGlobalFilter
       luaFilters <- peek (-1)
       runAll luaFilters pd
+ where
+  addMediaBagModule = do
+    Lua.getglobal "pandoc"
+    push "mediabag"
+    pushMediaBagModule mbRef
+    Lua.rawset (-3)
+  registerFormat = do
+    push format
+    Lua.setglobal "FORMAT"
+
 
 pushGlobalFilter :: Lua ()
 pushGlobalFilter = do
diff --git a/src/Text/Pandoc/Lua/PandocModule.hs b/src/Text/Pandoc/Lua/PandocModule.hs
index afb9aeca6..ffd681d30 100644
--- a/src/Text/Pandoc/Lua/PandocModule.hs
+++ b/src/Text/Pandoc/Lua/PandocModule.hs
@@ -15,6 +15,10 @@ You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 -}
+{-# LANGUAGE CPP #-}
+#if !MIN_VERSION_hslua(0,9,0)
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+#endif
 {- |
    Module      : Text.Pandoc.Lua.PandocModule
    Copyright   : Copyright © 2017 Albert Krewinkel
@@ -25,28 +29,37 @@ Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 Pandoc module for lua.
 -}
-module Text.Pandoc.Lua.PandocModule ( pushPandocModule ) where
+module Text.Pandoc.Lua.PandocModule
+  ( pushPandocModule
+  , pushMediaBagModule
+  ) where
 
-import Control.Monad (unless)
+import Control.Monad (unless, zipWithM_)
 import Data.ByteString.Char8 (unpack)
 import Data.Default (Default (..))
+import Data.IORef
 import Data.Text (pack)
-import Foreign.Lua (Lua, Status (OK), NumResults, call, loadstring, liftIO,
-                    push, pushHaskellFunction, rawset)
-import Text.Pandoc.Class (readDataFile, runIO, runIOorExplode, setUserDataDir)
+import Foreign.Lua (Lua, FromLuaStack, ToLuaStack, NumResults, liftIO)
+import Text.Pandoc.Class (fetchMediaResource, readDataFile, runIO,
+                          runIOorExplode, setUserDataDir)
 import Text.Pandoc.Options (ReaderOptions(readerExtensions))
 import Text.Pandoc.Lua.StackInstances ()
 import Text.Pandoc.Readers (Reader (..), getReader)
+import Text.Pandoc.MIME (MimeType)
+
+import qualified Foreign.Lua as Lua
+import qualified Data.ByteString.Lazy as BL
+import qualified Text.Pandoc.MediaBag as MB
 
 -- | Push the "pandoc" on the lua stack.
 pushPandocModule :: Maybe FilePath -> Lua ()
 pushPandocModule datadir = do
   script <- liftIO (pandocModuleScript datadir)
-  status <- loadstring script
-  unless (status /= OK) $ call 0 1
-  push "__read"
-  pushHaskellFunction readDoc
-  rawset (-3)
+  status <- Lua.loadstring script
+  unless (status /= Lua.OK) $ Lua.call 0 1
+  Lua.push "__read"
+  Lua.pushHaskellFunction readDoc
+  Lua.rawset (-3)
 
 -- | Get the string representation of the pandoc module
 pandocModuleScript :: Maybe FilePath -> IO String
@@ -56,14 +69,98 @@ pandocModuleScript datadir = unpack <$>
 readDoc :: String -> String -> Lua NumResults
 readDoc formatSpec content = do
   case getReader formatSpec of
-    Left  s      -> push s -- Unknown reader
+    Left  s      -> Lua.push s -- Unknown reader
     Right (reader, es) ->
       case reader of
         TextReader r -> do
           res <- liftIO $ runIO $ r def{ readerExtensions = es } (pack content)
           case res of
-            Left s   -> push $ show s -- error while reading
-            Right pd -> push pd       -- success, push Pandoc
-        _  -> push "Only string formats are supported at the moment."
+            Left s   -> Lua.push $ show s -- error while reading
+            Right pd -> Lua.push pd       -- success, push Pandoc
+        _  -> Lua.push "Only string formats are supported at the moment."
+  return 1
+
+--
+-- MediaBag submodule
+--
+pushMediaBagModule :: IORef MB.MediaBag -> Lua ()
+pushMediaBagModule mediaBagRef = do
+  Lua.newtable
+  addFunction "insert" (insertMediaFn mediaBagRef)
+  addFunction "lookup" (lookupMediaFn mediaBagRef)
+  addFunction "list" (mediaDirectoryFn mediaBagRef)
+  addFunction "fetch" (insertResource mediaBagRef)
+  return ()
+ where
+  addFunction name fn = do
+    Lua.push name
+    Lua.pushHaskellFunction fn
+    Lua.rawset (-3)
+
+insertMediaFn :: IORef MB.MediaBag
+              -> FilePath
+              -> OrNil MimeType
+              -> BL.ByteString
+              -> Lua NumResults
+insertMediaFn mbRef fp nilOrMime contents = do
+  liftIO . modifyIORef' mbRef $ MB.insertMedia fp (toMaybe nilOrMime) contents
+  return 0
+
+lookupMediaFn :: IORef MB.MediaBag
+              -> FilePath
+              -> Lua NumResults
+lookupMediaFn mbRef fp = do
+  res <- MB.lookupMedia fp <$> liftIO (readIORef mbRef)
+  case res of
+    Nothing -> Lua.pushnil *> return 1
+    Just (mimeType, contents) -> do
+      Lua.push mimeType
+      Lua.push contents
+      return 2
+
+mediaDirectoryFn :: IORef MB.MediaBag
+                 -> Lua NumResults
+mediaDirectoryFn mbRef = do
+  dirContents <- MB.mediaDirectory <$> liftIO (readIORef mbRef)
+  Lua.newtable
+  zipWithM_ addEntry [1..] dirContents
   return 1
+ where
+  addEntry :: Int -> (FilePath, MimeType, Int) -> Lua ()
+  addEntry idx (fp, mimeType, contentLength) = do
+    Lua.newtable
+    Lua.push "path" *> Lua.push fp *> Lua.rawset (-3)
+    Lua.push "type" *> Lua.push mimeType *> Lua.rawset (-3)
+    Lua.push "length" *> Lua.push contentLength *> Lua.rawset (-3)
+    Lua.rawseti (-2) idx
+
+insertResource :: IORef MB.MediaBag
+               -> String
+               -> OrNil String
+               -> Lua NumResults
+insertResource mbRef src sourceUrlOrNil = do
+  (fp, mimeType, bs) <- liftIO . runIOorExplode $
+                        fetchMediaResource (toMaybe sourceUrlOrNil) src
+  liftIO $ print (fp, mimeType)
+  insertMediaFn mbRef fp (OrNil mimeType) bs
+
+--
+-- Helper types and orphan instances
+--
+
+newtype OrNil a = OrNil { toMaybe :: Maybe a }
+
+instance FromLuaStack a => FromLuaStack (OrNil a) where
+  peek idx = do
+    noValue <- Lua.isnil idx
+    if noValue
+      then return (OrNil Nothing)
+      else OrNil . Just <$> Lua.peek idx
+
+#if !MIN_VERSION_hslua(0,9,0)
+instance ToLuaStack BL.ByteString where
+  push = Lua.push . BL.toStrict
 
+instance FromLuaStack BL.ByteString where
+  peek = fmap BL.fromStrict . Lua.peek
+#endif
diff --git a/test/Tests/Lua.hs b/test/Tests/Lua.hs
index fea813890..ba6196ccb 100644
--- a/test/Tests/Lua.hs
+++ b/test/Tests/Lua.hs
@@ -7,10 +7,11 @@ import Test.Tasty (TestTree, localOption)
 import Test.Tasty.HUnit (Assertion, assertEqual, testCase)
 import Test.Tasty.QuickCheck (ioProperty, testProperty, QuickCheckTests(..))
 import Text.Pandoc.Arbitrary ()
-import Text.Pandoc.Definition (Block, Inline, Meta, Pandoc)
 import Text.Pandoc.Builder ( (<>), bulletList, doc, doubleQuoted, emph
                            , linebreak, rawBlock, singleQuoted, para, plain
                            , space, str, strong)
+import Text.Pandoc.Class (runIOorExplode)
+import Text.Pandoc.Definition (Block, Inline, Meta, Pandoc)
 import Text.Pandoc.Lua
 
 import Foreign.Lua
@@ -80,8 +81,11 @@ tests = map (localOption (QuickCheckTests 20))
 
 assertFilterConversion :: String -> FilePath -> Pandoc -> Pandoc -> Assertion
 assertFilterConversion msg filterPath docIn docExpected = do
-  docRes <- runLuaFilter (Just "../data") ("lua" </> filterPath) [] docIn
-  assertEqual msg docExpected docRes
+  docEither <- runIOorExplode $
+               runLuaFilter (Just "../data") ("lua" </> filterPath) [] docIn
+  case docEither of
+    Left _ -> fail "lua filter failed"
+    Right docRes -> assertEqual msg docExpected docRes
 
 roundtripEqual :: (Eq a, FromLuaStack a, ToLuaStack a) => a -> IO Bool
 roundtripEqual x = (x ==) <$> roundtripped