Repository: PrefectHQ/prefect

Issue:
Allow exporter arguments in Jupyter ExecuteNotebook task
## Current behavior

When running the `jupyter.jupyter.ExecuteNotebook` task with `output_format='html'` the default settings for the HTMLExporter are used. There is no way to pass arguments to this exporter.

## Proposed behavior

Allow passing arguments to the HTMLExporter.

## Implementation suggestion

Something like `html_exporter = nbconvert.HTMLExporter(**exporter_kwargs)` on the following line:
https://github.com/PrefectHQ/prefect/blob/master/src/prefect/tasks/jupyter/jupyter.py#L65

## Example usecase

This allows you to exclude code cells, only showing their output, in the exported html document by passing the `exclude_input=True` argument to the exporter.


Discussion:
Since this would only apply to the html exporter and not the json output perhaps it'd make sense to allow make `output_format` a `Union[str, Exporter]` so an initialized exporter can be passed in directly instead with whatever arguments you want (base `Exporter` type at https://github.com/jupyter/nbconvert/blob/master/nbconvert/exporters/exporter.py#L49)

Resolving patch:
diff --git a/src/prefect/tasks/jupyter/jupyter.py b/src/prefect/tasks/jupyter/jupyter.py
--- a/src/prefect/tasks/jupyter/jupyter.py
+++ b/src/prefect/tasks/jupyter/jupyter.py
@@ -18,8 +18,12 @@ class ExecuteNotebook(Task):
             Can also be provided post-initialization by calling this task instance
         - parameters (dict, optional): dictionary of parameters to use for the notebook
             Can also be provided at runtime
-        - output_format (str, optional): Notebook output format.
-            Currently supported: json, html (default: json)
+        - output_format (str, optional): Notebook output format, should be a valid
+            nbconvert Exporter name. 'json' is treated as 'notebook'.
+            Valid exporter names: asciidoc, custom, html, latex, markdown,
+            notebook, pdf, python, rst, script, slides, webpdf. (default: notebook)
+        - exporter_kwargs (dict, optional): The arguments used for initializing
+            the exporter.
         - kernel_name (string, optional): kernel name to run the notebook with.
             If not provided, the default kernel will be used.
         - **kwargs: additional keyword arguments to pass to the Task constructor
@@ -29,7 +33,8 @@ def __init__(
         self,
         path: str = None,
         parameters: dict = None,
-        output_format: str = "json",
+        output_format: str = "notebook",
+        exporter_kwargs: dict = None,
         kernel_name: str = None,
         **kwargs
     ):
@@ -37,33 +42,40 @@ def __init__(
         self.parameters = parameters
         self.output_format = output_format
         self.kernel_name = kernel_name
+        self.exporter_kwargs = exporter_kwargs
         super().__init__(**kwargs)
 
-    @defaults_from_attrs("path", "parameters", "output_format")
+    @defaults_from_attrs("path", "parameters", "output_format", "exporter_kwargs")
     def run(
         self,
         path: str = None,
         parameters: dict = None,
         output_format: str = None,
+        exporter_kwargs: dict = None,
     ) -> str:
         """
-        Run a Jupyter notebook and output as HTML or JSON
+        Run a Jupyter notebook and output as HTML, notebook, or other formats.
 
         Args:
         - path (string, optional): path to fetch the notebook from; can also be
             a cloud storage path
         - parameters (dict, optional): dictionary of parameters to use for the notebook
-        - output_format (str, optional): Notebook output format.
-            Currently supported: json, html (default: json)
+        - output_format (str, optional): Notebook output format, should be a valid
+            nbconvert Exporter name. 'json' is treated as 'notebook'.
+            Valid exporter names: asciidoc, custom, html, latex, markdown,
+            notebook, pdf, python, rst, script, slides, webpdf. (default: notebook)
+        - exporter_kwargs (dict, optional): The arguments used for initializing
+            the exporter.
         """
         nb: nbformat.NotebookNode = pm.execute_notebook(
             path, "-", parameters=parameters, kernel_name=self.kernel_name
         )
         if output_format == "json":
-            return nbformat.writes(nb)
-        if output_format == "html":
-            html_exporter = nbconvert.HTMLExporter()
-            (body, resources) = html_exporter.from_notebook_node(nb)
-            return body
+            output_format = "notebook"
 
-        raise NotImplementedError("Notebook output %s not supported", output_format)
+        if exporter_kwargs is None:
+            exporter_kwargs = {}
+
+        exporter = nbconvert.get_exporter(output_format)
+        body, resources = nbconvert.export(exporter, nb, **exporter_kwargs)
+        return body


Test patch:
diff --git a/tests/tasks/jupyter/sample_notebook.ipynb b/tests/tasks/jupyter/sample_notebook.ipynb
--- a/tests/tasks/jupyter/sample_notebook.ipynb
+++ b/tests/tasks/jupyter/sample_notebook.ipynb
@@ -22,6 +22,15 @@
    "source": [
     "print(f\"a*b={a*b}\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This cell contains commented Python code"
+   ]
   }
  ],
  "metadata": {
diff --git a/tests/tasks/jupyter/test_jupyter.py b/tests/tasks/jupyter/test_jupyter.py
--- a/tests/tasks/jupyter/test_jupyter.py
+++ b/tests/tasks/jupyter/test_jupyter.py
@@ -11,6 +11,19 @@ def test_jupyter_html_output():
     )
     output = task.run()
     assert "a*b=10" in output
+    assert "# This cell contains commented Python code" in output
+
+
+def test_jupyter_html_output_with_exporter_kwargs():
+    task = ExecuteNotebook(
+        "tests/tasks/jupyter/sample_notebook.ipynb",
+        parameters=dict(a=5),
+        output_format="html",
+        exporter_kwargs={"exclude_input": True},
+    )
+    output = task.run()
+    assert "a*b=10" in output
+    assert "# This cell contains commented Python code" not in output
 
 
 def test_jupyter_json_output():
@@ -22,3 +35,14 @@ def test_jupyter_json_output():
     output = task.run()
     _ = json.loads(output)  # try loading the JSON string
     assert "a*b=10" in output
+
+
+def test_jupyter_notebook_output():
+    task = ExecuteNotebook(
+        "tests/tasks/jupyter/sample_notebook.ipynb",
+        parameters=dict(a=5),
+        output_format="notebook",
+    )
+    output = task.run()
+    _ = json.loads(output)  # try loading the JSON string
+    assert "a*b=10" in output