Pyplot.hs

{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE Unsafe     #-}
{-|
Module      : Text.Pandoc.Filter.Pyplot
Description : Pandoc filter to create Matplotlib figures from code blocks
Copyright   : (c) Laurent P René de Cotret, 2018
License     : MIT
Maintainer  : laurent.decotret@outlook.com
Stability   : stable
Portability : portable

This module defines a Pandoc filter @makePlot@ that can be 
used to walk over a Pandoc document and generate figures from
Python code blocks.
-}
module Text.Pandoc.Filter.Pyplot (
        makePlot
      , makePlot'
      , plotTransform
      , PandocPyplotError(..)
      , showError
    ) where

import           Control.Monad                  ((>=>))
import qualified Data.Map.Strict                as M
import           System.Directory               (doesDirectoryExist)
import           System.FilePath                (isValid, replaceExtension, takeDirectory)

import           Text.Pandoc.Definition
import           Text.Pandoc.Walk               (walkM)

import           Text.Pandoc.Filter.Scripting   

-- | Possible errors returned by the filter
data PandocPyplotError = ScriptError Int                -- ^ Running Python script has yielded an error
                       | InvalidTargetError FilePath    -- ^ Invalid figure path
                       | MissingDirectoryError FilePath -- ^ Directory where to save figure does not exist
                       | BlockingCallError              -- ^ Python script contains a block call to 'show()'

-- | Datatype containing all parameters required
-- to run pandoc-pyplot
data FigureSpec = FigureSpec 
    { target  :: FilePath     -- ^ filepath where generated figure will be saved
    , alt     :: String       -- ^ Alternate text for the figure (optional)
    , caption :: String       -- ^ Figure caption (optional)
    , script  :: PythonScript -- ^ Source code for the figure
    , blockAttrs :: Attr      -- ^ Attributes not related to @pandoc-pyplot@ will be propagated
    }

-- | Get the source code for a script including provisions to capture
-- the output.
scriptWithCapture :: FigureSpec -> PythonScript
scriptWithCapture spec = addPlotCapture (target spec) (script spec)

-- | Determine where to save the script source based on plot target
scriptSourcePath :: FigureSpec -> FilePath
scriptSourcePath spec = replaceExtension (target spec) ".txt"

-- | Get the source code for a figure script in a presentable way
presentableScript :: FigureSpec -> PythonScript
presentableScript spec = mconcat [ "# Source code for ", target spec, "\n", script spec ]

-- Keys that pandoc-pyplot will look for in code blocks
targetKey, altTextKey, captionKey :: String
targetKey  = "plot_target"
altTextKey = "plot_alt"
captionKey = "plot_caption"

-- | Determine inclusion specifications from Block attributes.
-- Note that the target key is required, but all other parameters are optional
parseFigureSpec :: Block -> Maybe FigureSpec
parseFigureSpec (CodeBlock (id', cls, attrs) content) = 
    createInclusion <$> M.lookup targetKey attrs'
    where
        attrs' = M.fromList attrs
        inclusionKeys = [ targetKey, altTextKey, captionKey ]
        filteredAttrs = filter (\(k,_) -> k `notElem` inclusionKeys) attrs
        createInclusion fname = FigureSpec
            { target     = fname
            , alt        = M.findWithDefault "Figure generated by pandoc-pyplot" altTextKey attrs'
            , caption    = M.findWithDefault mempty captionKey attrs'
            , script     = content
            -- Propagate attributes that are not related to pandoc-pyplot
            , blockAttrs = (id', cls, filteredAttrs)
            }
parseFigureSpec _ = Nothing

-- | Main routine to include Matplotlib plots.
-- Code blocks containing the attributes @plot_target@ are considered
-- Python plotting scripts. All other possible blocks are ignored.
-- The source code is also saved in another file, which can be access by 
-- clicking the image
makePlot' :: Block -> IO (Either PandocPyplotError Block)
makePlot' block = 
    case parseFigureSpec block of
        -- Could not parse - leave code block unchanged
        Nothing -> return $ Right block
        -- Could parse : run the script and capture output
        Just spec -> do
            let figurePath = target spec
                figureDir = takeDirectory figurePath
                scriptSource = script spec
            
            -- Check that the directory in which to save the figure exists
            validDirectory <- doesDirectoryExist figureDir
            
            if | not (isValid figurePath)         -> return $ Left $ InvalidTargetError figurePath
               | not validDirectory               -> return $ Left $ MissingDirectoryError figureDir
               | hasBlockingShowCall scriptSource -> return $ Left $ BlockingCallError
               | otherwise -> do 
                    
                -- Running the script happens on the next line
                -- Note that the script is slightly modified to be able to capture the output
                result <- runTempPythonScript (scriptWithCapture spec)
                
                case result of
                    ScriptFailure code -> return $ Left $ ScriptError code
                    ScriptSuccess -> do 
                        -- Save the original script into a separate file
                        -- so it can be inspected
                        -- Note : using a .txt file allows to view source directly
                        --        in the browser, in the case of HTML output
                        let sourcePath = scriptSourcePath spec
                        writeFile sourcePath (presentableScript spec)
                        
                        -- Propagate attributes that are not related to pandoc-pyplot
                        let relevantAttrs = blockAttrs spec
                            image     = Image relevantAttrs [Str $ alt spec] (figurePath, "")
                            srcTarget = (sourcePath, "Click on this figure to see the source code")

                        -- TODO: use FigureSpec caption
                        -- We make the figure be a link to the source code
                        return $ Right $ Para [
                            Link nullAttr [image] srcTarget
                            ]

-- | Translate filter error to an error message
showError :: PandocPyplotError -> String
showError (ScriptError exitcode)          = "Script error: plot could not be generated. Exit code " <> (show exitcode)
showError (InvalidTargetError fname)      = "Target filename " <> fname <> " is not valid."
showError (MissingDirectoryError dirname) = "Target directory " <> dirname <> " does not exist." 
showError BlockingCallError               = "Script contains a blocking call to show, like 'plt.show()'"

-- | Highest-level function that can be walked over a Pandoc tree.
-- All code blocks that have the 'plot_target' parameter will be considered
-- figures.
makePlot :: Block -> IO Block
makePlot = makePlot' >=> either (fail . showError) return

-- | Walk over an entire Pandoc document, changing appropriate code blocks
-- into figures.
plotTransform :: Pandoc -> IO Pandoc
plotTransform = walkM makePlot