ipynb-0.3: test/roundtrip.hs

{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad (unless)
import Data.Aeson (Value (..), eitherDecode, encode, object, (.=))
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as B
import Data.Ipynb
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Lazy.Encoding as TLE
import Lens.Micro
import Lens.Micro.Aeson
import System.Directory
import System.FilePath
import Test.Tasty
import Test.Tasty.HUnit
import Data.Monoid

main :: IO ()
main = do
  let rtdir = "test" </> "rt-files"
  createDirectoryIfMissing False rtdir
  fs <- map (rtdir </>) . filter isIpynb <$> getDirectoryContents rtdir
  defaultMain $ testGroup "tests"
    [ testGroup "round-trip tests" $ map rtTest fs
    , testGroup "mime data tests" mimeTests
    , testCase "execute_result with null execution_count" $ do
        let inRaw =
              "{\"nbformat\":4,\"nbformat_minor\":5,\"metadata\":{}," <>
              "\"cells\":[{\"cell_type\":\"code\",\"metadata\":{}," <>
              "\"source\":\"\",\"execution_count\":null,\"outputs\":" <>
              "[{\"output_type\":\"execute_result\",\"execution_count\":null," <>
              "\"data\":{\"text/plain\":\"1\"},\"metadata\":{}}]}]}"
        (nb :: Notebook NbV4) <- either error return $ eitherDecode inRaw
        case map cellType (notebookCells nb) of
          [Code{ codeOutputs = [ExecuteResult{ executeCount = ec }] }] ->
            ec @?= Nothing
          _ -> assertFailure "unexpected notebook structure"
        (nb' :: Notebook NbV4) <- either error return $ eitherDecode (encode nb)
        nb' @?= nb
    ]

-- Parse a v4 notebook with a single code cell whose only output
-- is display_data with the given data object, returning the MimeBundle.
parseV4Data :: BL.ByteString -> Either String MimeBundle
parseV4Data dat = do
  (nb :: Notebook NbV4) <- eitherDecode $
    "{\"nbformat\":4,\"nbformat_minor\":5,\"metadata\":{},\"cells\":" <>
    "[{\"cell_type\":\"code\",\"metadata\":{},\"source\":\"\"," <>
    "\"execution_count\":1,\"outputs\":[{\"output_type\":\"display_data\"," <>
    "\"metadata\":{},\"data\":" <> dat <> "}]}]}"
  case map cellType (notebookCells nb) of
    [Code{ codeOutputs = [DisplayData{ displayData = d }] }] -> Right d
    _ -> Left "unexpected notebook structure"

mimeTests :: [TestTree]
mimeTests =
  [ testCase "textual data that looks like base64 stays textual" $
      parseV4Data "{\"application/javascript\":\"test\"}" @?=
        Right (MimeBundle (M.fromList
          [("application/javascript", TextualData "test")]))
  , testCase "image data is decoded from base64" $
      parseV4Data "{\"image/png\":\"dGVzdA==\\n\"}" @?=
        Right (MimeBundle (M.fromList [("image/png", BinaryData "test")]))
  , testCase "svg data stays textual" $
      parseV4Data "{\"image/svg+xml\":[\"<svg>\\n\",\"</svg>\"]}" @?=
        Right (MimeBundle (M.fromList
          [("image/svg+xml", TextualData "<svg>\n</svg>")]))
  , testCase "non-textual data under a non-json mime type is kept as json" $
      parseV4Data "{\"application/vnd.foo\":{\"a\":1}}" @?=
        Right (MimeBundle (M.fromList
          [("application/vnd.foo", JsonData (object ["a" .= (1 :: Int)]))]))
  , testCase "v3 svg, json, and pdf data is read" $
      parseV3Data
        "\"svg\":\"<svg/>\",\"json\":{\"a\":1},\"pdf\":\"dGVzdA==\"" @?=
        Right (MimeBundle (M.fromList
          [ ("image/svg+xml", TextualData "<svg/>")
          , ("application/json", JsonData (object ["a" .= (1 :: Int)]))
          , ("application/pdf", BinaryData "test") ]))
  , testCase "v3 mime data survives write and read" $ do
      let bundle = MimeBundle (M.fromList
            [ ("text/plain", TextualData "hi")
            , ("image/svg+xml", TextualData "<svg/>")
            , ("application/json", JsonData (object ["a" .= (1 :: Int)]))
            , ("application/pdf", BinaryData "test")
            , ("application/vnd.foo", TextualData "foo") ])
          nb :: Notebook NbV3
          nb = Notebook mempty (3, 0)
                 [Cell (Code (Just 1) [DisplayData bundle mempty])
                       Nothing (Source ["x"]) mempty Nothing]
      fmap (map cellType . notebookCells)
        (eitherDecode (encode nb) :: Either String (Notebook NbV3)) @?=
        Right (map cellType (notebookCells nb))
  ]

-- Like parseV4Data, but for a v3 notebook; the argument is a
-- comma-separated sequence of JSON fields to add to the display_data
-- output.
parseV3Data :: BL.ByteString -> Either String MimeBundle
parseV3Data fields = do
  (nb :: Notebook NbV3) <- eitherDecode $
    "{\"nbformat\":3,\"nbformat_minor\":0,\"metadata\":{},\"worksheets\":" <>
    "[{\"cells\":[{\"cell_type\":\"code\",\"input\":\"\",\"metadata\":{}," <>
    "\"outputs\":[{\"output_type\":\"display_data\",\"metadata\":{}," <>
    fields <> "}]}]}]}"
  case map cellType (notebookCells nb) of
    [Code{ codeOutputs = [DisplayData{ displayData = d }] }] -> Right d
    _ -> Left "unexpected notebook structure"

isIpynb :: FilePath -> Bool
isIpynb fp = takeExtension fp == ".ipynb"

rtTest :: FilePath -> TestTree
rtTest fp = testCase fp $ do
  inRaw <- BL.readFile fp
  let format = inRaw ^? key "nbformat" . _Number
  case format of
    Just 4 -> rtTest4 inRaw
    _      -> rtTest3 inRaw

rtTest3 :: BL.ByteString -> IO ()
rtTest3 inRaw = do
  (nb :: Notebook NbV3) <- either error return $ eitherDecode inRaw
  let outRaw = encode nb
  (nb' :: Notebook NbV3) <- either error return $ eitherDecode outRaw
  -- test that (write . read) == id
  assertEqual "write . read != read" nb nb'

rtTest4 :: BL.ByteString -> IO ()
rtTest4 inRaw = do
  (nb :: Notebook NbV4) <- either error return $ eitherDecode inRaw
  let outRaw = encode nb
  (nb' :: Notebook NbV4) <- either error return $ eitherDecode outRaw
  -- test that (write . read) == id
  assertEqual "write . read != read" nb nb'