Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add flag --draft-notebook for generating a skeleton notebook for manual editing (e.g. in VSCode). #1284

Merged
merged 4 commits into from Nov 29, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
13 changes: 13 additions & 0 deletions docs/snakefiles/rules.rst
Expand Up @@ -945,6 +945,19 @@ with
The last dependency is advisable in order to enable autoformatting of notebook cells when editing.
When using other languages than Python in the notebook, one needs to additionally add the respective kernel, e.g. ``r-irkernel`` for R support.

When using an IDE with built-in Jupyter support, an alternative to ``--edit-notebook`` is ``--draft-notebook``.
Instead of firing up a notebook server, ``--draft-notebook`` just creates a skeleton notebook for editing within the IDE.
In addition, it prints instructions for configuring the IDE's notebook environment to use the interpreter from the
Conda environment defined in the corresponding rule.
For example, running

.. code-block:: console

snakemake --cores 1 --edit-notebook test.txt --use-conda

will generate skeleton code in ``notebooks/hello.py.ipynb`` and additionally print instructions on how to open and execute the notebook in VSCode.


Protected and Temporary Files
-----------------------------

Expand Down
18 changes: 15 additions & 3 deletions snakemake/__init__.py
Expand Up @@ -299,7 +299,7 @@ def snakemake(
export_cwl (str): Compile workflow to CWL and save to given file
log_handler (function): redirect snakemake output to this custom log handler, a function that takes a log message dictionary (see below) as its only argument (default None). The log message dictionary for the log handler has to following entries:
keep_incomplete (bool): keep incomplete output files of failed jobs
edit_notebook (object): "notebook.Listen" object to configuring notebook server for interactive editing of a rule notebook. If None, do not edit.
edit_notebook (object): "notebook.EditMode" object to configuring notebook server for interactive editing of a rule notebook. If None, do not edit.
scheduler (str): Select scheduling algorithm (default ilp)
scheduler_ilp_solver (str): Set solver for ilp scheduler.
overwrite_groups (dict): Rule to group assignments (default None)
Expand Down Expand Up @@ -1489,6 +1489,13 @@ def get_argument_parser(profile=None):

group_notebooks = parser.add_argument_group("NOTEBOOKS")

group_notebooks.add_argument(
"--draft-notebook",
metavar="TARGET",
help="Draft a skeleton notebook for the rule used to generate the given target file. This notebook "
"can then be opened in a jupyter server, exeucted and implemented until ready. After saving, it "
"will automatically be reused in non-interactive mode by Snakemake for subsequent jobs.",
)
group_notebooks.add_argument(
"--edit-notebook",
metavar="TARGET",
Expand Down Expand Up @@ -2721,12 +2728,17 @@ def open_browser():
)
log_handler.append(wms_logger.log_handler)

if args.edit_notebook:
if args.draft_notebook:
from snakemake import notebook

args.target = [args.draft_notebook]
args.edit_notebook = notebook.EditMode(draft_only=True)
elif args.edit_notebook:
from snakemake import notebook

args.target = [args.edit_notebook]
args.force = True
args.edit_notebook = notebook.Listen(args.notebook_listen)
args.edit_notebook = notebook.EditMode(args.notebook_listen)

aggregated_wait_for_files = args.wait_for_files
if args.wait_for_files_file is not None:
Expand Down
2 changes: 1 addition & 1 deletion snakemake/executors/__init__.py
Expand Up @@ -512,7 +512,7 @@ def job_args_and_prepare(self, job):
self.workflow.cleanup_scripts,
job.shadow_dir,
job.jobid,
self.workflow.edit_notebook,
self.workflow.edit_notebook if self.dag.is_edit_notebook_job(job) else None,
self.workflow.conda_base_path,
job.rule.basedir,
self.workflow.sourcecache.runtime_cache_path,
Expand Down
4 changes: 4 additions & 0 deletions snakemake/jobs.py
Expand Up @@ -1004,6 +1004,10 @@ def postprocess(
latency_wait=None,
keep_metadata=True,
):
if self.dag.is_edit_notebook_job(self):
# No postprocessing necessary, we have just created the skeleton notebook and
# execution will anyway stop afterwards.
return
if assume_shared_fs:
if not error and handle_touch:
self.dag.handle_touch(self)
Expand Down
61 changes: 54 additions & 7 deletions snakemake/notebook.py
@@ -1,4 +1,6 @@
from abc import abstractmethod
import os, sys
from pathlib import Path
from urllib.error import URLError
import tempfile
import re
Expand All @@ -10,15 +12,17 @@
from snakemake.logging import logger
from snakemake.common import is_local_file
from snakemake.common import ON_WINDOWS
from snakemake.sourcecache import SourceCache
from snakemake.sourcecache import SourceCache, infer_source_file

KERNEL_STARTED_RE = re.compile(r"Kernel started: (?P<kernel_id>\S+)")
KERNEL_SHUTDOWN_RE = re.compile(r"Kernel shutdown: (?P<kernel_id>\S+)")


class Listen:
def __init__(self, arg):
self.ip, self.port = arg.split(":")
class EditMode:
def __init__(self, server_addr=None, draft_only=False):
if server_addr is not None:
self.ip, self.port = server_addr.split(":")
self.draft_only = draft_only


def get_cell_sources(source):
Expand All @@ -33,20 +37,24 @@ class JupyterNotebook(ScriptBase):

editable = True

def draft(self, listen):
def draft(self):
import nbformat

preamble = self.get_preamble()
nb = nbformat.v4.new_notebook()
self.insert_preamble_cell(preamble, nb)

nb["cells"].append(nbformat.v4.new_code_cell("# start coding here"))
nb["metadata"] = {"language_info": {"name": self.get_language_name()}}

os.makedirs(os.path.dirname(self.local_path), exist_ok=True)

with open(self.local_path, "wb") as out:
out.write(nbformat.writes(nb).encode())

def draft_and_edit(self, listen):
self.draft()

self.source = open(self.local_path).read()

self.evaluate(edit=listen)
Expand All @@ -73,6 +81,7 @@ def execute_script(self, fname, edit=None):

with tempfile.TemporaryDirectory() as tmp:
if edit is not None:
assert not edit.draft_only
logger.info("Opening notebook for editing.")
cmd = (
"jupyter notebook --browser ':' --no-browser --log-level ERROR --ip {edit.ip} --port {edit.port} "
Expand Down Expand Up @@ -132,6 +141,14 @@ def remove_preamble_cell(self, notebook):
# remove old preamble
del notebook["cells"][preamble]

@abstractmethod
def get_language_name(self):
...

@abstractmethod
def get_interpreter_exec(self):
...


class PythonJupyterNotebook(JupyterNotebook):
def get_preamble(self):
Expand Down Expand Up @@ -163,6 +180,12 @@ def get_preamble(self):
preamble_addendum=preamble_addendum,
)

def get_language_name(self):
return "python"

def get_interpreter_exec(self):
return "python"


class RJupyterNotebook(JupyterNotebook):
def get_preamble(self):
Expand Down Expand Up @@ -193,6 +216,12 @@ def get_preamble(self):
preamble_addendum=preamble_addendum,
)

def get_language_name(self):
return "r"

def get_interpreter_exec(self):
return "RScript"


def get_exec_class(language):
exec_class = {
Expand Down Expand Up @@ -266,6 +295,7 @@ def notebook(
else:
source = None
is_local = True
path = infer_source_file(path)

exec_class = get_exec_class(language)

Expand Down Expand Up @@ -295,7 +325,24 @@ def notebook(
is_local,
)

if draft:
executor.draft(listen=edit)
if edit is None:
executor.evaluate(edit=edit)
elif edit.draft_only:
executor.draft()
msg = "Generated skeleton notebook:\n{} ".format(path)
if conda_env and not container_img:
msg += (
"\n\nEditing with VSCode:\nOpen notebook, run command 'Select notebook kernel' (Ctrl+Shift+P or Cmd+Shift+P), and choose:"
"\n{}\n".format(
str(Path(conda_env) / "bin" / executor.get_interpreter_exec())
)
)
msg += (
"\nEditing with Jupyter CLI:"
"\nconda activate {}\njupyter notebook {}\n".format(conda_env, path)
)
logger.info(msg)
elif draft:
executor.draft_and_edit(listen=edit)
else:
executor.evaluate(edit=edit)
46 changes: 46 additions & 0 deletions tests/test_jupyter_notebook_draft/Snakefile
@@ -0,0 +1,46 @@
shell.executable("bash")

rule all:
input:
'result_final.txt',
'book.result_final.txt',

rule foo:
output:
fname = 'data.txt'
run:
with open(output.fname, 'w') as fd:
fd.write('result of serious computation')

rule bar:
input:
infile = 'data.txt'
output:
outfile = 'result_intermediate.txt'
conda:
'env.yaml'
notebook:
'Notebook.py.ipynb'

rule baz:
input:
infile = 'result_intermediate.txt'
output:
outfile = 'result_final.txt'
log:
notebook = 'Notebook_Processed.ipynb'
conda:
'env.yaml'
notebook:
'Notebook.py.ipynb'


rule wild:
input:
infile = 'result_intermediate.txt'
output:
outfile = '{what}.result_final.txt'
conda:
'env.yaml'
notebook:
'Note{wildcards.what}.py.ipynb'
1 change: 1 addition & 0 deletions tests/test_jupyter_notebook_draft/data.txt
@@ -0,0 +1 @@
result of serious computation
7 changes: 7 additions & 0 deletions tests/test_jupyter_notebook_draft/env.yaml
@@ -0,0 +1,7 @@
channels:
- conda-forge
- bioconda
dependencies:
- python >=3.5
- jupyter
- ipykernel
@@ -0,0 +1,38 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "f8c5fca4",
"metadata": {
"tags": [
"snakemake-job-properties"
]
},
"outputs": [],
"source": [
"\n",
"######## snakemake preamble start (automatically inserted, do not edit) ########\n",
"import sys; sys.path.extend(['/Users/johannes/scms/snakemake', '/Users/johannes/scms/snakemake/tests/test_jupyter_notebook_draft']); import pickle; snakemake = pickle.loads(b'\\x80\\x03csnakemake.script\\nSnakemake\\nq\\x00)\\x81q\\x01}q\\x02(X\\x05\\x00\\x00\\x00inputq\\x03csnakemake.io\\nInputFiles\\nq\\x04)\\x81q\\x05X\\x08\\x00\\x00\\x00data.txtq\\x06a}q\\x07(X\\x06\\x00\\x00\\x00_namesq\\x08}q\\tX\\x06\\x00\\x00\\x00infileq\\nK\\x00N\\x86q\\x0bsX\\x12\\x00\\x00\\x00_allowed_overridesq\\x0c]q\\r(X\\x05\\x00\\x00\\x00indexq\\x0eX\\x04\\x00\\x00\\x00sortq\\x0feh\\x0ecfunctools\\npartial\\nq\\x10cbuiltins\\ngetattr\\nq\\x11csnakemake.io\\nNamedlist\\nq\\x12X\\x0f\\x00\\x00\\x00_used_attributeq\\x13\\x86q\\x14Rq\\x15\\x85q\\x16Rq\\x17(h\\x15)}q\\x18X\\x05\\x00\\x00\\x00_nameq\\x19h\\x0esNtq\\x1abh\\x0fh\\x10h\\x15\\x85q\\x1bRq\\x1c(h\\x15)}q\\x1dh\\x19h\\x0fsNtq\\x1ebh\\nh\\x06ubX\\x06\\x00\\x00\\x00outputq\\x1fcsnakemake.io\\nOutputFiles\\nq )\\x81q!X\\x17\\x00\\x00\\x00result_intermediate.txtq\"a}q#(h\\x08}q$X\\x07\\x00\\x00\\x00outfileq%K\\x00N\\x86q&sh\\x0c]q\\'(h\\x0eh\\x0feh\\x0eh\\x10h\\x15\\x85q(Rq)(h\\x15)}q*h\\x19h\\x0esNtq+bh\\x0fh\\x10h\\x15\\x85q,Rq-(h\\x15)}q.h\\x19h\\x0fsNtq/bh%h\"ubX\\x06\\x00\\x00\\x00paramsq0csnakemake.io\\nParams\\nq1)\\x81q2}q3(h\\x08}q4h\\x0c]q5(h\\x0eh\\x0feh\\x0eh\\x10h\\x15\\x85q6Rq7(h\\x15)}q8h\\x19h\\x0esNtq9bh\\x0fh\\x10h\\x15\\x85q:Rq;(h\\x15)}q<h\\x19h\\x0fsNtq=bubX\\t\\x00\\x00\\x00wildcardsq>csnakemake.io\\nWildcards\\nq?)\\x81q@}qA(h\\x08}qBh\\x0c]qC(h\\x0eh\\x0feh\\x0eh\\x10h\\x15\\x85qDRqE(h\\x15)}qFh\\x19h\\x0esNtqGbh\\x0fh\\x10h\\x15\\x85qHRqI(h\\x15)}qJh\\x19h\\x0fsNtqKbubX\\x07\\x00\\x00\\x00threadsqLK\\x01X\\t\\x00\\x00\\x00resourcesqMcsnakemake.io\\nResources\\nqN)\\x81qO(K\\x01K\\x01X0\\x00\\x00\\x00/var/folders/l0/9bhq7fc12lgfknlx5gyxckv00000gp/TqPe}qQ(h\\x08}qR(X\\x06\\x00\\x00\\x00_coresqSK\\x00N\\x86qTX\\x06\\x00\\x00\\x00_nodesqUK\\x01N\\x86qVX\\x06\\x00\\x00\\x00tmpdirqWK\\x02N\\x86qXuh\\x0c]qY(h\\x0eh\\x0feh\\x0eh\\x10h\\x15\\x85qZRq[(h\\x15)}q\\\\h\\x19h\\x0esNtq]bh\\x0fh\\x10h\\x15\\x85q^Rq_(h\\x15)}q`h\\x19h\\x0fsNtqabhSK\\x01hUK\\x01hWhPubX\\x03\\x00\\x00\\x00logqbcsnakemake.io\\nLog\\nqc)\\x81qd}qe(h\\x08}qfh\\x0c]qg(h\\x0eh\\x0feh\\x0eh\\x10h\\x15\\x85qhRqi(h\\x15)}qjh\\x19h\\x0esNtqkbh\\x0fh\\x10h\\x15\\x85qlRqm(h\\x15)}qnh\\x19h\\x0fsNtqobubX\\x06\\x00\\x00\\x00configqp}qqX\\x04\\x00\\x00\\x00ruleqrX\\x03\\x00\\x00\\x00barqsX\\x0f\\x00\\x00\\x00bench_iterationqtNX\\t\\x00\\x00\\x00scriptdirquX@\\x00\\x00\\x00/Users/johannes/scms/snakemake/tests/test_jupyter_notebook_draftqvub.'); from snakemake.logging import logger; logger.printshellcmds = False; import os; os.chdir(r'/Users/johannes/scms/snakemake/tests/test_jupyter_notebook_draft');\n",
"######## snakemake preamble end #########\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4df7dcd5",
"metadata": {},
"outputs": [],
"source": [
"# start coding here"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
12 changes: 12 additions & 0 deletions tests/tests.py
Expand Up @@ -1172,6 +1172,18 @@ def test_jupyter_notebook():
run(dpath("test_jupyter_notebook"), use_conda=True)


def test_jupyter_notebook_draft():
from snakemake.notebook import EditMode

run(
dpath("test_jupyter_notebook_draft"),
use_conda=True,
edit_notebook=EditMode(draft_only=True),
targets=["result_intermediate.txt"],
check_md5=False,
)


def test_github_issue456():
run(dpath("test_github_issue456"))

Expand Down