Merge pull request #12 from JuliaPluto/screenshot-cells

JuliaPluto · Feb 12, 2024 · e19c26d · e19c26d
2 parents a620310 + 6bf048a
commit e19c26d
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 40 deletions.
diff --git a/node/bin.js b/node/bin.js
@@ -3,22 +3,16 @@ import { pdf } from "./export.js"
 import path from "path"
 import fileUrl from "file-url"
 
-const fileInput = process.argv[2]
-const fileOutput = process.argv[3]
+const html_path_arg = process.argv[2]
+const pdf_path_arg = process.argv[3]
 const options = JSON.parse(process.argv[4])
+const screenshot_path_arg = process.argv[5]
+const screenshot_options = JSON.parse(process.argv[6])
 
-if (!fileInput) {
-    console.error("ERROR: First program argument must be a Pluto notebook path or URL")
-    process.exit(1)
-}
-if (!fileOutput) {
-    console.error("ERROR: Second program argument must be the PDF output path")
-    process.exit(1)
-}
+const input_url = html_path_arg.startsWith("http://") || html_path_arg.startsWith("https://") ? html_path_arg : fileUrl(path.resolve(html_path_arg))
+const pdf_path = path.resolve(pdf_path_arg)
+const screenshot_path = screenshot_path_arg == "" ? null : path.resolve(screenshot_path_arg)
 
-const exportUrl = fileInput.startsWith("http://") || fileInput.startsWith("https://") ? fileInput : fileUrl(path.resolve(fileInput))
-const pdf_path = path.resolve(fileOutput)
-
-await pdf(exportUrl, pdf_path, options)
+await pdf(input_url, pdf_path, options, screenshot_path, screenshot_options)
 
 process.exit()
diff --git a/node/export.js b/node/export.js
@@ -1,5 +1,6 @@
 import p from "puppeteer"
 import chalk from "chalk"
+import path from "path"
 
 function sleep(time) {
     return new Promise((resolve, reject) => {
@@ -9,7 +10,7 @@ function sleep(time) {
     })
 }
 
-export async function pdf(url, pdf_path, options, beforeClose = async () => {}) {
+export async function pdf(url, pdf_path, options, screenshot_dir, screenshot_options, { beforeClose = async () => {} } = {}) {
     const browser = await p.launch()
     console.log("Initiated headless browser")
     const page = await browser.newPage()
@@ -24,34 +25,66 @@ export async function pdf(url, pdf_path, options, beforeClose = async () => {})
         height: 1000,
     })
 
-    while (true) {
-        const queued = await page.evaluate(`Array.from(document.getElementsByClassName('queued')).map(x => x.id)`)
-        const running = await page.evaluate(`Array.from(document.getElementsByClassName('running')).map(x => x.id)`)
-        const cells = await page.evaluate(`Array.from(document.getElementsByTagName('pluto-cell')).map(x => x.id)`)
-        const bodyClasses = await page.evaluate(`document.body.getAttribute('class')`)
-
-        if (running.length > 0) {
-            process.stdout.write(`\rRunning cell ${chalk.yellow(`${cells.length - queued.length}/${cells.length}`)} ${chalk.cyan(`[${running[0]}]`)}`)
-        }
-
-        if (!(bodyClasses.includes("loading") || queued.length > 0 || running.length > 0)) {
-            process.stdout.write(`\rRunning cell ${chalk.yellow(`${cells.length}/${cells.length}`)}`)
-            console.log()
-            break
-        }
-
-        await sleep(250)
-    }
+    await waitForPlutoBusy(page, false, { timeout: 30 * 1000 })
 
     console.log("Exporting as pdf...")
     await page.pdf({
         path: pdf_path,
         ...options,
     })
+    if (screenshot_dir != null) {
+        await screenshot_cells(page, screenshot_dir, screenshot_options)
+    }
 
     console.log(chalk.green("Exported ✓") + " ... cleaning up")
 
     await beforeClose()
-
     await browser.close()
 }
+
+/**
+ * @param {p.Page} page
+ * @param {string} screenshot_dir
+ */
+async function screenshot_cells(page, screenshot_dir, { outputOnly, scale }) {
+    const cells = /** @type {String[]} */ (await page.evaluate(`Array.from(document.querySelectorAll('pluto-cell')).map(x => x.id)`))
+
+    for (let cell_id of cells) {
+        const cell = await page.$(`[id="${cell_id}"]${outputOnly ? " > pluto-output" : ""}`)
+        if (cell) {
+            await cell.scrollIntoView()
+            const rect = await cell.boundingBox()
+            if (rect == null) {
+                throw new Error(`Cell ${cell_id} is not visible`)
+            }
+            const imgpath = path.join(screenshot_dir, `${cell_id}.png`)
+
+            await cell.screenshot({ path: imgpath, clip: { ...rect, scale }, omitBackground: false })
+            console.log(`Screenshot ${cell_id} saved to ${imgpath}`)
+        }
+    }
+}
+
+const timeout = (delay) =>
+    new Promise((r) => {
+        setTimeout(r, delay)
+    })
+
+const waitForPlutoBusy = async (page, iWantBusiness, options) => {
+    await timeout(1000)
+    await page.waitForFunction(
+        (iWantBusiness) => {
+            let quiet = //@ts-ignore
+                (document?.body?._update_is_ongoing ?? false) === false &&
+                //@ts-ignore
+                (document?.body?._js_init_set?.size ?? 0) === 0 &&
+                document?.body?.classList?.contains("loading") === false &&
+                document?.querySelector(`pluto-cell.running, pluto-cell.queued, pluto-cell.internal_test_queued`) == null
+
+            return iWantBusiness ? !quiet : quiet
+        },
+        options,
+        iWantBusiness
+    )
+    await timeout(1000)
+}
diff --git a/src/PlutoPDF.jl b/src/PlutoPDF.jl
@@ -36,19 +36,34 @@ const default_options = (
     displayHeaderFooter=false,
 )
 
-function html_to_pdf(html_path::AbstractString, output_path::Union{AbstractString,Nothing}=nothing; 
-    options=default_options, 
+const screenshot_default_options = (
+    outputOnly=false,
+    scale=2,
+)
+
+function html_to_pdf(
+    html_path::AbstractString, 
+    output_path::Union{AbstractString,Nothing}=nothing,
+    screenshot_dir_path::Union{AbstractString,Nothing}=nothing; 
+    options=default_options,
+    screenshot_options=screenshot_default_options,
     open=true, 
     console_output=true
 )
     bin_script = normpath(joinpath(@__DIR__, "../node/bin.js"))
 
     output_path = tamepath(something(output_path, Pluto.numbered_until_new(splitext(html_path)[1]; suffix=".pdf", create_file=false)))
+
+    screenshot_dir_path = if screenshot_dir_path === nothing
+        nothing
+    else
+        mkpath(tamepath(screenshot_dir_path))
+    end
 
     @info "Generating pdf..."
     cmd = `$(node()) $bin_script $(tamepath(html_path)) $(output_path) $(JSON.json(
         (; default_options..., options...)
-    ))`
+    )) $(something(screenshot_dir_path, "")) $(JSON.json((; screenshot_default_options..., screenshot_options...)))`
     if console_output
         run(cmd)
     else
@@ -77,7 +92,12 @@ Run a notebook, generate an Export HTML and then print it to a PDF file!
 # Options
 The `options` keyword argument can be a named tuple to configure the PDF export. The possible options can be seen in the [docs for `puppeteer.PDFOptions`](https://pptr.dev/api/puppeteer.pdfoptions). You don't need to specify all options, for example: `options=(format="A5",)` will work.
 """
-function pluto_to_pdf(notebook_path::AbstractString, output_path::Union{AbstractString,Nothing}=nothing; kwargs...)
+function pluto_to_pdf(
+    notebook_path::AbstractString,
+    output_path::Union{AbstractString,Nothing}=nothing,
+    screenshot_dir_path::Union{AbstractString,Nothing}=nothing; 
+    kwargs...
+)
     c = Pluto.Configuration.from_flat_kwargs(;
         disable_writing_notebook_files = true,
         lazy_workspace_creation = true,
@@ -93,7 +113,7 @@ function pluto_to_pdf(notebook_path::AbstractString, output_path::Union{Abstract
 
     output_path = something(output_path, Pluto.numbered_until_new(Pluto.without_pluto_file_extension(notebook_path); suffix=".pdf", create_file=false))
 
-    html_to_pdf(filename, output_path; kwargs...)
+    html_to_pdf(filename, output_path, screenshot_dir_path; kwargs...)
 end
 
 function __init__()

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -6,16 +6,38 @@ testfile = download(
     "https://raw.githubusercontent.com/fonsp/Pluto.jl/main/sample/Tower%20of%20Hanoi.jl",
 )
 
+testfile2 = download("https://raw.githubusercontent.com/JuliaPluto/PlutoSliderServer.jl/08dafcca4073551cb4192d442dc3e8c33123b952/test/dir1/a.jl")
+
 
 is_CI = get(ENV, "CI", "no") == "no"
 
-outfile = pluto_to_pdf(testfile; open=is_CI, options=(format="A5",))
+outfile = tempname(; cleanup=false) * ".pdf"
+outdir = tempname(; cleanup=false)
+
+@info "Files" outfile outdir testfile testfile2
 
+result = pluto_to_pdf(testfile, outfile, outdir; open=is_CI, options=(format="A5",))
+
+@test result == outfile
 
 @test isfile(outfile)
 @test dirname(outfile) == dirname(testfile)
 @test endswith(outfile, ".pdf")
 
+@test isdir(outdir)
+filez = readdir(outdir)
+@test length(filez) == 28
+@test all(endswith.(filez, ".png"))
+@test length(read(joinpath(outdir, filez[1]))) > 1000
+
+
+
+outfile2 = pluto_to_pdf(testfile2; open=is_CI, options=(format="A5",))
+@info "Result" outfile2
+@test isfile(outfile2)
+@test dirname(outfile2) == dirname(testfile2)
+@test endswith(outfile2, ".pdf")
+
 output_dir = get(ENV, "TEST_OUTPUT_DIR", nothing)
 
 if @show(output_dir) isa String