wrap in a function

sosiristseng · Apr 3, 2024 · 509f27c · 509f27c
1 parent 25c73b0
commit 509f27c
Showing 1 changed file with 73 additions and 68 deletions.
diff --git a/run.jl b/run.jl
@@ -9,92 +9,97 @@ using IJulia
     Pkg.activate(Base.current_project())
 end
 
-basedir = get(ENV, "DOCDIR", "docs") # Defaults to docs/
-cachedir = get(ENV, "NBCACHE", ".cache") # Defaults to .cache/
-ipynbs = String[]
-litnbs = String[]
+function main(; basedir = get(ENV, "DOCDIR", "docs"), cachedir = get(ENV, "NBCACHE", ".cache"), printtable=true)
+    ipynbs = String[]
+    litnbs = String[]
 
-# Collect the list of notebooks
-for (root, dirs, files) in walkdir(basedir)
-    for file in files
-        if endswith(file, ".ipynb") || endswith(file, ".jl")
-            nb = joinpath(root, file)
-            shaval = read(nb, String) |> sha256 |> bytes2hex
-            @info "Notebook $(nb): hash=$(shaval)"
-            shafilename = joinpath(cachedir, root, splitext(file)[1] * ".sha")
-            # Cache hit
-            if isfile(shafilename) && read(shafilename, String) == shaval
-                @info "Notebook $(nb) cache hits and will not be executed."
-            # Cache miss
-            else
-                @info "Notebook $(nb) cache misses. Writing hash to $(shafilename)."
-                mkpath(dirname(shafilename))
-                write(shafilename, shaval)
-                if endswith(file, ".ipynb")
-                    push!(ipynbs, nb)
-                elseif endswith(file, ".jl")
-                    push!(litnbs, nb)
+    mkpath(cachedir)
+
+    # Remove cached notebook and sha files if there is no corresponding notebook
+    for (root, dirs, files) in walkdir(cachedir)
+        for file in files
+            if endswith(file, ".ipynb") || endswith(file, ".sha")
+                fn = joinpath(joinpath(splitpath(root)[2:end]), splitext(file)[1])
+                nb = fn * ".ipynb"
+                lit = fn * ".jl"
+                if !isfile(nb) && !isfile(lit)
+                    fullfn = joinpath(root, file)
+                    @info "Notebook $(nb) or $(lit) not found. Removing $(fullfn)."
+                    rm(fullfn)
                 end
             end
         end
     end
-end
 
-# Remove cached notebook and sha files if there is no corresponding notebook
-for (root, dirs, files) in walkdir(cachedir)
-    for file in files
-        if endswith(file, ".ipynb") || endswith(file, ".sha")
-            fn = joinpath(joinpath(splitpath(root)[2:end]), splitext(file)[1])
-            nb = fn * ".ipynb"
-            lit = fn * ".jl"
-            if !isfile(nb) && !isfile(lit)
-                fullfn = joinpath(root, file)
-                @info "Notebook $(nb) or $(lit) not found. Removing $(fullfn)."
-                rm(fullfn)
+    # Collect the list of notebooks
+    for (root, dirs, files) in walkdir(basedir)
+        for file in files
+            if endswith(file, ".ipynb") || endswith(file, ".jl")
+                nb = joinpath(root, file)
+                shaval = read(nb, String) |> sha256 |> bytes2hex
+                @info "Notebook $(nb): hash=$(shaval)"
+                shafilename = joinpath(cachedir, root, splitext(file)[1] * ".sha")
+                # Cache hit
+                if isfile(shafilename) && read(shafilename, String) == shaval
+                    @info "Notebook $(nb) cache hits and will not be executed."
+                # Cache miss
+                else
+                    @info "Notebook $(nb) cache misses. Writing hash to $(shafilename)."
+                    mkpath(dirname(shafilename))
+                    write(shafilename, shaval)
+                    if endswith(file, ".ipynb")
+                        push!(ipynbs, nb)
+                    elseif endswith(file, ".jl")
+                        push!(litnbs, nb)
+                    end
+                end
             end
         end
     end
-end
 
-# Execute literate notebooks in worker process(es)
-ts_lit = pmap(litnbs; on_error=ex->NaN) do nb
-    outdir = joinpath(cachedir, dirname(nb))
-    @elapsed Literate.notebook(nb, outdir; mdstrings=true)
-end
+    # Execute literate notebooks in worker process(es)
+    ts_lit = pmap(litnbs; on_error=ex->NaN) do nb
+        outdir = joinpath(cachedir, dirname(nb))
+        @elapsed Literate.notebook(nb, outdir; mdstrings=true)
+    end
 
-# Remove worker processes in Distributed.jl
-rmprocs(workers())
+    # Remove worker processes to release some memory
+    rmprocs(workers())
 
-# Debug notebooks one by one if there are errors
-for (nb, t) in zip(litnbs, ts_lit)
-    if isnan(t)
-        println("Debugging notebook: ", nb)
-        try
-            withenv("JULIA_DEBUG" => "Literate") do
-                Literate.notebook(nb, dirname(nb); mdstrings=true)
+    # Debug notebooks one by one if there are errors
+    for (nb, t) in zip(litnbs, ts_lit)
+        if isnan(t)
+            println("Debugging notebook: ", nb)
+            try
+                withenv("JULIA_DEBUG" => "Literate") do
+                    Literate.notebook(nb, dirname(nb); mdstrings=true)
+                end
+            catch e
+                println(e)
             end
-        catch e
-            println(e)
         end
     end
-end
 
-any(isnan, ts_lit) && error("Please check literate notebook error(s).")
+    any(isnan, ts_lit) && error("Please check literate notebook error(s).")
 
-# Install IJulia kernel
-IJulia.installkernel("Julia", "--project=@.", "--heap-size-hint=3G")
+    # Install IJulia kernel
+    IJulia.installkernel("Julia", "--project=@.", "--heap-size-hint=3G")
 
-# nbconvert command array
-ntasks = parse(Int, get(ENV, "NBCONVERT_JOBS", "1"))
-kernelname = "--ExecutePreprocessor.kernel_name=julia-1.$(VERSION.minor)"
-execute = ifelse(get(ENV, "ALLOWERRORS", " ") == "true", "--execute --allow-errors", "--execute")
-timeout = "--ExecutePreprocessor.timeout=" * get(ENV, "TIMEOUT", "-1")
-cmds = [`jupyter nbconvert --to notebook $(execute) $(timeout) $(kernelname) --output $(joinpath(abspath(pwd()), cachedir, nb)) $(nb)` for nb in ipynbs]
+    # nbconvert command array
+    ntasks = parse(Int, get(ENV, "NBCONVERT_JOBS", "1"))
+    kernelname = "--ExecutePreprocessor.kernel_name=julia-1.$(VERSION.minor)"
+    execute = ifelse(get(ENV, "ALLOWERRORS", " ") == "true", "--execute --allow-errors", "--execute")
+    timeout = "--ExecutePreprocessor.timeout=" * get(ENV, "TIMEOUT", "-1")
+    cmds = [`jupyter nbconvert --to notebook $(execute) $(timeout) $(kernelname) --output $(joinpath(abspath(pwd()), cachedir, nb)) $(nb)` for nb in ipynbs]
+
+    # Run the nbconvert commands in parallel
+    ts_ipynb = asyncmap(cmds; ntasks) do cmd
+        @elapsed run(cmd)
+    end
 
-# Run the nbconvert commands in parallel
-ts_ipynb = asyncmap(cmds; ntasks) do cmd
-    @elapsed run(cmd)
+    # Print execution result
+    printtable && pretty_table([litnbs ts_lit; ipynbs ts_ipynb], header=["Notebook", "Elapsed (s)"])
 end
 
-# Print execution result
-pretty_table([litnbs ts_lit; ipynbs ts_ipynb], header=["Notebook", "Elapsed (s)"])
+# Run code
+main()