Implement expand extension interface and Cell struct

src/QuartoNotebookWorker/src/QuartoNotebookWorker.jl

@@ -1,5 +1,11 @@
 module QuartoNotebookWorker
 
+# Exports:
+
+export Cell
+export expand
+
+
 walk(x, _, outer) = outer(x)
 walk(x::Expr, inner, outer) = outer(Expr(x.head, map(inner, x.args)...))
 postwalk(f, x) = walk(x, x -> postwalk(f, x), f)
@@ -83,6 +89,7 @@ include("InlineDisplay.jl")
 include("NotebookState.jl")
 include("NotebookInclude.jl")
 include("refresh.jl")
+include("cell_expansion.jl")
 include("render.jl")
 include("utilities.jl")
 include("ojs_define.jl")

src/QuartoNotebookWorker/src/cell_expansion.jl

@@ -0,0 +1,86 @@
+"""
+    struct Cell
+    Cell(content; code = nothing, options = Dict())
+
+`content` is either a callable object, which should have a 0-argument method
+which, when called, represents the evaluation of the cell, or any other
+non-callable object, which will simply be treated as the output value of the
+cell. Should you wish to use a callable object as the output of a cell, rather
+than calling it, you can wrap it in a `Returns` object.
+
+`code` is the mock source code of the cell, and is not parsed or evaluated, but
+will be rendered in the final output generated by `quarto`. When `code` is not
+provided then the code block is hidden with `echo: false`. `options` represents
+the cell options for the mock cell and will impact the rendering of the final
+output.
+"""
+struct Cell
+    thunk::Base.Callable
+    code::String
+    options::Dict{String,Any}
+
+    function Cell(object; code::Union{String,Nothing} = nothing, options::Dict = Dict())
+        thunk = _make_thunk(object)
+        if isnothing(code)
+            options["echo"] = false
+            code = ""
+        end
+        new(thunk, code, options)
+    end
+end
+
+_make_thunk(c::Base.Callable) = c
+# Returns is only available on 1.7 and upwards.
+_make_thunk(other) = @static @isdefined(Returns) ? Returns(other) : () -> other
+
+"""
+    expand(object::T) -> Vector{Cell}
+
+Define the vector of `Cell`s that an object of type `T` should expand to. See
+`Cell` documentation for more details on `Cell` creation.
+    
+This function is meant to be extended by 3rd party packages using Julia's
+package extension mechanism. For example, within the `Project.toml` of a package
+called `ExamplePackage` add the following:
+
+```toml
+[weakdeps]
+QuartoNotebookWorker = "38328d9c-a911-4051-bc06-3f7f556ffeda"
+
+[extensions]
+ExamplePackageQuartoNotebookWorkerExt = "QuartoNotebookWorker"
+```
+
+Then create a file `ext/ExamplePackageQuartoNotebookWorkerExt.jl` with the the contents
+
+```julia
+module ExamplePackageQuartoNotebookWorkerExt
+
+import ExamplePackage
+import QuartoNotebookWorker
+
+function QuartoNotebookWorker.expand(obj::ExamplePackage.ExampleType)
+    return [
+        QuartoNotebookWorker.Cell("This is the cell result."; code = "# Mock code goes here."),
+        QuartoNotebookWorker.Cell("This is a second cell."),
+    ]
+end
+
+end
+```
+"""
+expand(@nospecialize(_)) = nothing
+
+struct CellExpansionError <: Exception
+    message::String
+end
+
+_is_expanded(@nospecialize(_), ::Nothing) = false
+_is_expanded(@nospecialize(_), ::Vector{Cell}) = true
+function _is_expanded(@nospecialize(original), @nospecialize(result))
+    throw(
+        CellExpansionError(
+            "invalid cell expansion result for `expand(::$(typeof(original)))`. Expected the result to`Vector{Cell}`, got `$(typeof(result))`.",
+        ),
+    )
+end

src/QuartoNotebookWorker/src/render.jl

@@ -4,9 +4,16 @@ function render(
     line::Integer,
     cell_options::AbstractDict = Dict{String,Any}(),
 )
-    return Base.@invokelatest(
+    # This records whether the outermost cell is an expandable cell, which we
+    # then return to the server so that it can decide whether to treat the cell
+    # results it gets back as an expansion or not. We can't decide this
+    # statically since expansion depends on whether the runtime type of the cell
+    # output is `is_expandable` or not. Recursive calls to `_render_thunk` don't
+    # matter to the server, it's just the outermost cell that matters.
+    is_expansion_ref = Ref(false)
+    result = Base.@invokelatest(
         collect(
-            _render_thunk(code, cell_options) do
+            _render_thunk(code, cell_options, is_expansion_ref) do
                 Base.@invokelatest include_str(
                     NotebookState.notebook_module(),
                     code;
@@ -16,6 +23,7 @@ function render(
             end,
         )
     )
+    return (result, is_expansion_ref[])
 end
 
 # Recursively render cell thunks. This might be an `include_str` call,
@@ -26,133 +34,53 @@ function _render_thunk(
     thunk::Base.Callable,
     code::AbstractString,
     cell_options::AbstractDict = Dict{String,Any}(),
+    is_expansion_ref::Ref{Bool} = Ref(false),
 )
     captured, display_results = with_inline_display(thunk, cell_options)
-    if get(cell_options, "expand", false) === true
-        if captured.error
+
+    # Attempt to expand the cell. This requires the cell result to have a method
+    # defined for the `QuartoNotebookWorker.expand` function. We only attempt to
+    # run expansion if the cell didn't error. Cell expansion can itself error,
+    # so we need to catch that and return an error cell if that's the case.
+    expansion = nothing
+    is_expansion = false
+    if !captured.error
+        try
+            expansion = expand(captured.value)
+            is_expansion = _is_expanded(captured.value, expansion)
+        catch error
+            backtrace = catch_backtrace()
             return ((;
                 code = "", # an expanded cell that errored can't have returned code
                 cell_options = Dict{String,Any}(), # or options
                 results = Dict{String,@NamedTuple{error::Bool, data::Vector{UInt8}}}(),
                 display_results,
                 output = captured.output,
-                error = string(typeof(captured.value)),
+                error = string(typeof(error)),
                 backtrace = collect(
-                    eachline(
-                        IOBuffer(
-                            clean_bt_str(
-                                captured.error,
-                                captured.backtrace,
-                                captured.value,
-                            ),
-                        ),
-                    ),
+                    eachline(IOBuffer(clean_bt_str(true, backtrace, error))),
                 ),
             ),)
-        else
-            function invalid_return_value_cell(
-                errmsg;
-                code = "",
-                cell_options = Dict{String,Any}(),
-            )
-                return ((;
-                    code,
-                    cell_options,
-                    results = Dict{String,@NamedTuple{error::Bool, data::Vector{UInt8}}}(),
-                    display_results,
-                    output = captured.output,
-                    error = "Invalid return value for expanded cell",
-                    backtrace = collect(eachline(IOBuffer(errmsg))),
-                ),)
-            end
+        end
+        # Track in this side-channel whether the cell is an expansion or not.
+        is_expansion_ref[] = is_expansion
+    end
 
-            if !(Base.@invokelatest Base.isiterable(typeof(captured.value)))
-                return invalid_return_value_cell(
-                    """
-                    Return value of a cell with `expand: true` is not iterable.
-                    The returned value must iterate objects that each have a `thunk`
-                    property which contains a function that returns the cell output.
-                    Instead, the returned value was:
-                    $(repr(captured.value))
-                    """,
+    if is_expansion
+        # A cell expansion with `expand` might itself also contain
+        # cells that expand to multiple cells, so we need to flatten
+        # the results to a single list of cells before passing back
+        # to the server. Cell expansion is recursive.
+        return _flatmap(expansion) do cell
+            wrapped = function ()
+                return QuartoNotebookWorker.Packages.IOCapture.capture(
+                    cell.thunk;
+                    rethrow = InterruptException,
+                    color = true,
                 )
             end
-
-            # A cell expansion with `expand` might itself also contain
-            # cells that expand to multiple cells, so we need to flatten
-            # the results to a single list of cells before passing back
-            # to the server. Cell expansion is recursive.
-            return _flatmap(enumerate(captured.value)) do (i, cell)
-
-                code = _getproperty(cell, :code, "")
-                options = _getproperty(Dict{String,Any}, cell, :options)
-
-                if !(code isa String)
-                    return invalid_return_value_cell(
-                        """
-                        While iterating over the elements of the return value of a cell with
-                        `expand: true`, a value was found at position $i which has a `code` property
-                        that is not of the expected type `String`. The value was:
-                        $(repr(cell.code))
-                        """,
-                    )
-                end
-
-                if !(options isa Dict{String})
-                    return invalid_return_value_cell(
-                        """
-                        While iterating over the elements of the return value of a cell with
-                        `expand: true`, a value was found at position $i which has a `options` property
-                        that is not of the expected type `Dict{String}`. The value was:
-                        $(repr(cell.options))
-                        """;
-                        code,
-                    )
-                end
-
-                if !hasproperty(cell, :thunk)
-                    return invalid_return_value_cell(
-                        """
-                        While iterating over the elements of the return value of a cell with
-                        `expand: true`, a value was found at position $i which does not have a
-                        `thunk` property. Every object in the iterator returned from an expanded
-                        cell must have a property `thunk` with a function that returns
-                        the output of the cell.
-                        The object without a `thunk` property was:
-                        $(repr(cell))
-                        """;
-                        code,
-                        cell_options = options,
-                    )
-                end
-
-                if !(cell.thunk isa Base.Callable)
-                    return invalid_return_value_cell(
-                        """
-                        While iterating over the elements of the return value of a cell with
-                        `expand: true` a value was found at position $i which has a `thunk`
-                        property that is not a function of type `Base.Callable`.
-                        Every object in the iterator returned from an expanded
-                        cell must have a property `thunk` with a function that returns
-                        the output of the cell. Instead, the returned value was:
-                        $(repr(cell.thunk))
-                        """;
-                        code,
-                        cell_options = options,
-                    )
-                end
-
-                wrapped = function ()
-                    return QuartoNotebookWorker.Packages.IOCapture.capture(
-                        cell.thunk;
-                        rethrow = InterruptException,
-                        color = true,
-                    )
-                end
-
-                # **The recursive call:**
-                return Base.@invokelatest _render_thunk(wrapped, code, options)
-            end
+            # **The recursive call:**
+            return Base.@invokelatest _render_thunk(wrapped, cell.code, cell.options)
         end
     else
         results = Base.@invokelatest render_mimetypes(

src/server.jl

@@ -612,12 +612,9 @@ function evaluate_raw_cells!(
 
     @maybe_progress showprogress "$header" for (nth, chunk) in enumerate(chunks)
         if chunk.type === :code
-            expand_cell = get(chunk.cell_options, "expand", false) === true
-            # When we're not evaluating the code, or when there is an `expand`
-            # cell output then we immediately splice in the cell code. The
-            # results of evaluating an `expand` cell are added later on and are
-            # not considered direct outputs of this cell.
-            if !chunk.evaluate || expand_cell
+            if !chunk.evaluate
+                # Cells that are not evaluated are not executed, but they are
+                # still included in the notebook.
                 push!(
                     cells,
                     (;
@@ -626,12 +623,10 @@ function evaluate_raw_cells!(
                         metadata = (;),
                         source = process_cell_source(chunk.source),
                         outputs = [],
-                        execution_count = chunk.evaluate ? 1 : 0,
+                        execution_count = 0,
                     ),
                 )
-            end
-
-            if chunk.evaluate
+            else
                 chunk_callback(ith_chunk_to_evaluate, chunks_to_evaluate, chunk)
                 ith_chunk_to_evaluate += 1
 
@@ -646,7 +641,26 @@ function evaluate_raw_cells!(
                     $(chunk.cell_options),
                 ))
 
-                for (mth, remote) in enumerate(Malt.remote_eval_fetch(f.worker, expr))
+                worker_results, expand_cell = Malt.remote_eval_fetch(f.worker, expr)
+
+                # When the result of the cell evaluation is a cell expansion
+                # then we insert the original cell contents before the expanded
+                # cells as a mock cell similar to if it has `eval: false` set.
+                if expand_cell
+                    push!(
+                        cells,
+                        (;
+                            id = string(nth),
+                            cell_type = chunk.type,
+                            metadata = (;),
+                            source = process_cell_source(chunk.source),
+                            outputs = [],
+                            execution_count = 1,
+                        ),
+                    )
+                end
+
+                for (mth, remote) in enumerate(worker_results)
                     outputs = []
                     processed = process_results(remote.results)
 
@@ -800,7 +814,10 @@ function evaluate_raw_cells!(
                             # There should only ever be a single result from an
                             # inline evaluation since you can't pass cell
                             # options and so `expand` will always be `false`.
-                            remote = only(Malt.remote_eval_fetch(f.worker, expr))
+                            worker_results, expand_cell =
+                                Malt.remote_eval_fetch(f.worker, expr)
+                            expand_cell && error("inline code cells cannot be expanded")
+                            remote = only(worker_results)
                             if !isnothing(remote.error)
                                 # file location is not straightforward to determine with inline literals, but just printing the (presumably short)
                                 # code back instead of a location should be quite helpful