feat: Add support for Nix Flakes through nix develop

snakemake/deployment/nix_flake.py

@@ -0,0 +1,9 @@
+class NixFlakeEnv:
+
+    """Nix Flake Env"""
+
+    def __init__(self, nix_flake):
+        self.nix_flake = nix_flake
+
+    def shellcmd(self, cmd):
+        return f"nix develop {self.nix_flake} --command bash -c '{cmd}'"

snakemake/executors/local.py

@@ -158,6 +158,7 @@ def job_args_and_prepare(self, job: JobExecutorInterface):
             benchmark_repeats,
             conda_env,
             container_img,
+            job.nix_flake,
             self.workflow.deployment_settings.apptainer_args,
             env_modules,
             DeploymentMethod.APPTAINER
@@ -299,6 +300,7 @@ def run_wrapper(
     benchmark_repeats,
     conda_env,
     container_img,
+    nix_flake,
     singularity_args,
     env_modules,
     use_singularity,
@@ -377,6 +379,7 @@ def run_wrapper(
                             rule,
                             conda_env,
                             container_img,
+                            nix_flake,
                             singularity_args,
                             use_singularity,
                             env_modules,
@@ -407,6 +410,7 @@ def run_wrapper(
                                 rule,
                                 conda_env,
                                 container_img,
+                                nix_flake,
                                 singularity_args,
                                 use_singularity,
                                 env_modules,
@@ -435,6 +439,7 @@ def run_wrapper(
                     rule,
                     conda_env,
                     container_img,
+                    nix_flake,
                     singularity_args,
                     use_singularity,
                     env_modules,

snakemake/jobs.py

@@ -475,6 +475,10 @@ def container_img(self):
             return self.dag.container_imgs[self.container_img_url]
         return None
 
+    @property
+    def nix_flake(self):
+        return self.rule.nix_flake
+
     @property
     def env_modules(self):
         return self.rule.env_modules

snakemake/parser.py

@@ -413,6 +413,10 @@ class Container(RuleKeywordState):
     pass
 
 
+class NixFlake(RuleKeywordState):
+    pass
+
+
 class Containerized(RuleKeywordState):
     pass
 
@@ -462,7 +466,7 @@ def start(self):
         yield "\n"
         yield (
             "def __rule_{rulename}(input, output, params, wildcards, threads, "
-            "resources, log, rule, conda_env, container_img, "
+            "resources, log, rule, conda_env, container_img, nix_flake, "
             "singularity_args, use_singularity, env_modules, bench_record, jobid, "
             "is_shell, bench_iteration, cleanup_scripts, shadow_dir, edit_notebook, "
             "conda_base_path, basedir, runtime_sourcecache_path, {rule_func_marker}=True):".format(
@@ -565,7 +569,7 @@ class Script(AbstractCmd):
     def args(self):
         yield (
             ", basedir, input, output, params, wildcards, threads, resources, log, "
-            "config, rule, conda_env, conda_base_path, container_img, singularity_args, env_modules, "
+            "config, rule, conda_env, conda_base_path, container_img, nix_flake, singularity_args, env_modules, "
             "bench_record, jobid, bench_iteration, cleanup_scripts, shadow_dir, runtime_sourcecache_path"
         )
 
@@ -577,7 +581,7 @@ class Notebook(Script):
     def args(self):
         yield (
             ", basedir, input, output, params, wildcards, threads, resources, log, "
-            "config, rule, conda_env, conda_base_path, container_img, singularity_args, env_modules, "
+            "config, rule, conda_env, conda_base_path, container_img, nix_flake, singularity_args, env_modules, "
             "bench_record, jobid, bench_iteration, cleanup_scripts, shadow_dir, "
             "edit_notebook, runtime_sourcecache_path"
         )
@@ -590,7 +594,7 @@ class Wrapper(Script):
     def args(self):
         yield (
             ", input, output, params, wildcards, threads, resources, log, "
-            "config, rule, conda_env, conda_base_path, container_img, singularity_args, env_modules, "
+            "config, rule, conda_env, conda_base_path, container_img, nix_flake, singularity_args, env_modules, "
             "bench_record, workflow.workflow_settings.wrapper_prefix, jobid, bench_iteration, "
             "cleanup_scripts, shadow_dir, runtime_sourcecache_path"
         )
@@ -630,6 +634,7 @@ def args(self):
     conda=Conda,
     singularity=Singularity,
     container=Container,
+    nix_flake=NixFlake,
     containerized=Containerized,
     envmodules=EnvModules,
     wildcard_constraints=WildcardConstraints,

snakemake/ruleinfo.py

@@ -25,6 +25,7 @@ def __init__(self, func=None):
         self.benchmark = None
         self.conda_env = None
         self.container_img = None
+        self.nix_flake = None
         self.is_containerized = False
         self.env_modules = None
         self.wildcard_constraints = None

snakemake/rules.py

@@ -96,6 +96,7 @@ def __init__(self, *args, lineno=None, snakefile=None):
             self._conda_env = None
             self._container_img = None
             self.is_containerized = False
+            self._nix_flake = None
             self.env_modules = None
             self._group = None
             self._wildcard_names = None
@@ -144,6 +145,7 @@ def __init__(self, *args, lineno=None, snakefile=None):
             self._conda_env = other._conda_env
             self._container_img = other._container_img
             self.is_containerized = other.is_containerized
+            self._nix_flake = other._nix_flake
             self.env_modules = other.env_modules
             self._group = other.group
             self._wildcard_names = (
@@ -408,6 +410,14 @@ def container_img(self):
     def container_img(self, container_img):
         self._container_img = container_img
 
+    @property
+    def nix_flake(self):
+        return self._nix_flake
+
+    @nix_flake.setter
+    def nix_flake(self, nix_flake):
+        self._nix_flake = nix_flake
+
     @property
     def input(self):
         return self._input

snakemake/shell.py

@@ -18,6 +18,7 @@
 from snakemake.logging import logger
 from snakemake.deployment import singularity
 from snakemake.deployment.conda import Conda
+from snakemake.deployment.nix_flake import NixFlakeEnv
 from snakemake.exceptions import WorkflowError
 
 
@@ -157,6 +158,7 @@ def __new__(
         conda_env = context.get("conda_env", None)
         conda_base_path = context.get("conda_base_path", None)
         container_img = context.get("container_img", None)
+        nix_flake = context.get("nix_flake", None)
         env_modules = context.get("env_modules", None)
         shadow_dir = context.get("shadow_dir", None)
         resources = context.get("resources", {})
@@ -183,6 +185,8 @@ def __new__(
                 cmd = Conda(
                     container_img=container_img, prefix_path=conda_base_path
                 ).shellcmd(conda_env, cmd)
+        if nix_flake:
+            cmd = NixFlakeEnv(nix_flake).shellcmd(cmd)
 
         tmpdir = None
         if len(cmd.replace("'", r"'\''")) + 2 > MAX_ARG_LEN:

snakemake/workflow.py

@@ -1536,6 +1536,8 @@ def decorate(ruleinfo):
                 or ruleinfo.shellcmd
                 or ruleinfo.notebook
             )
+            if ruleinfo.nix_flake:
+                rule.nix_flake = ruleinfo.nix_flake
             if ruleinfo.container_img:
                 if invalid_rule:
                     raise RuleException(
@@ -1716,6 +1718,13 @@ def decorate(ruleinfo):
 
         return decorate
 
+    def nixflake(self, nix_flake):
+        def decorate(ruleinfo):
+            ruleinfo.nix_flake = nix_flake
+            return ruleinfo
+
+        return decorate
+
     def containerized(self, container_img):
         def decorate(ruleinfo):
             ruleinfo.container_img = container_img

tests/test_nix_flake/Snakefile

@@ -0,0 +1,21 @@
+rule all:
+    input:
+        expand("test{i}.out2", i=range(3)),
+
+rule a:
+    output:
+        "test{i}.out",
+    nix_flake:
+        ".#pyshell"
+    shell:
+        "python3 --version > {output} || true"
+
+rule b:
+    input:
+        "test{i}.out",
+    output:
+        "test{i}.out2",
+    nix_flake:
+        ".#rshell"
+    shell:
+        "R --version > {output} || true"

tests/test_nix_flake/expected-results/test0.out

@@ -0,0 +1 @@
+Python 3.10.11

tests/test_nix_flake/expected-results/test1.out

@@ -0,0 +1 @@
+Python 3.10.11

tests/test_nix_flake/expected-results/test2.out

@@ -0,0 +1 @@
+Python 3.10.11

tests/test_nix_flake/flake.nix

@@ -0,0 +1,29 @@
+{
+  description = "A very basic flake";
+
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs/23.05";
+  };
+
+  outputs = { self, nixpkgs }:
+    let
+      system = "x86_64-linux";
+      pkgs = import nixpkgs { inherit system; };
+    in
+    {
+
+      devShells.${system} = {
+        pyshell = pkgs.mkShell {
+          buildInputs = with pkgs; [
+            python3
+          ];
+        };
+        rshell = pkgs.mkShell {
+          buildInputs = with pkgs; [
+            R
+          ];
+        };
+      };
+
+    };
+}