adapt script support to also work with non stdlib modules

snakemake/common/__init__.py

@@ -7,9 +7,11 @@
 import contextlib
 import itertools
 import math
+import operator
 import platform
 import hashlib
 import inspect
+import sys
 import threading
 import uuid
 import os
@@ -26,7 +28,6 @@
 
 MIN_PY_VERSION = (3, 9)
 DYNAMIC_FILL = "__snakemake_dynamic__"
-SNAKEMAKE_SEARCHPATH = str(Path(__file__).parent.parent.parent)
 UUID_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_URL, "https://snakemake.readthedocs.io")
 NOTHING_TO_BE_DONE_MSG = (
     "Nothing to be done (all requested files are present and up to date)."
@@ -39,6 +40,13 @@
 IO_PROP_LIMIT = 100
 
 
+def get_snakemake_searchpaths():
+    paths = [str(Path(__file__).parent.parent.parent)] + [
+        path for path in sys.path if path.endswith("site-packages")
+    ]
+    return list(unique_justseen(paths))
+
+
 def mb_to_mib(mb):
     return int(math.ceil(mb * 0.95367431640625))
 
@@ -261,3 +269,14 @@ async def async_lock(_lock: threading.Lock):
         yield  # the lock is held
     finally:
         _lock.release()
+
+
+def unique_justseen(iterable, key=None):
+    """
+    List unique elements, preserving order. Remember only the element just seen.
+
+    From https://docs.python.org/3/library/itertools.html#itertools-recipes
+    """
+    # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
+    # unique_justseen('ABBcCAD', str.lower) --> A B c A D
+    return map(next, map(operator.itemgetter(1), itertools.groupby(iterable, key)))

snakemake/deployment/singularity.py

@@ -3,15 +3,16 @@
 __email__ = "johannes.koester@uni-due.de"
 __license__ = "MIT"
 
+from pathlib import Path
 import subprocess
 import shutil
 import os
 import hashlib
 
 from snakemake.common import (
+    get_snakemake_searchpaths,
     is_local_file,
     parse_uri,
-    SNAKEMAKE_SEARCHPATH,
 )
 from snakemake.exceptions import WorkflowError
 from snakemake.logging import logger
@@ -21,6 +22,12 @@
 SNAKEMAKE_MOUNTPOINT = "/mnt/snakemake"
 
 
+def get_snakemake_searchpath_mountpoints():
+    paths = get_snakemake_searchpaths()
+    base = Path("/mnt/snakemake_searchpaths")
+    return [base / f"item_{i}" for i in range(len(paths))]
+
+
 class Image:
     def __init__(self, url, dag, is_containerized):
         if " " in url:
@@ -110,7 +117,12 @@ def shellcmd(
 
     if is_python_script:
         # mount host snakemake module into container
-        args += f" --bind {repr(SNAKEMAKE_SEARCHPATH)}:{repr(SNAKEMAKE_MOUNTPOINT)}"
+        args += " ".join(
+            f" --bind {repr(searchpath)}:{repr(mountpoint)}"
+            for searchpath, mountpoint in zip(
+                get_snakemake_searchpaths(), get_snakemake_searchpath_mountpoints()
+            )
+        )
 
     if container_workdir:
         args += f" --pwd {repr(container_workdir)}"

snakemake/script.py

@@ -34,8 +34,8 @@
 from snakemake.shell import shell
 from snakemake.common import (
     MIN_PY_VERSION,
-    SNAKEMAKE_SEARCHPATH,
     ON_WINDOWS,
+    get_snakemake_searchpaths,
 )
 from snakemake.deployment import singularity
 
@@ -543,16 +543,14 @@ def generate_preamble(
         # Obtain search path for current snakemake module.
         # The module is needed for unpickling in the script.
         # We append it at the end (as a fallback).
-        snakemake_searchpath = SNAKEMAKE_SEARCHPATH
+        searchpaths = get_snakemake_searchpaths()
         if container_img is not None:
-            snakemake_searchpath = singularity.SNAKEMAKE_MOUNTPOINT
-        searchpaths = [snakemake_searchpath] + [
-            path for path in sys.path if path.endswith("site-packages")
-        ]
+            searchpaths = singularity.get_snakemake_searchpath_mountpoints()
 
         # Add the cache path to the search path so that other cached source files in the same dir
         # can be imported.
         if cache_path:
+            # TODO handle this in case of container_img, analogously to above
             cache_searchpath = os.path.dirname(cache_path)
             if cache_searchpath:
                 searchpaths.append(cache_searchpath)
@@ -1045,16 +1043,6 @@ def encode_namedlist(values):
 
         json_string = json.dumps(dict(snakemake))
 
-        # Obtain search path for current snakemake module.
-        # We append it at the end (as a fallback).
-        searchpath = SNAKEMAKE_SEARCHPATH
-        if container_img is not None:
-            searchpath = singularity.SNAKEMAKE_MOUNTPOINT
-        searchpath = repr(searchpath)
-        # For local scripts, add their location to the path in case they use path-based imports
-        if is_local:
-            searchpath += ", " + repr(path.get_basedir().get_path_or_uri())
-
         return textwrap.dedent(
             """
             json_typegen::json_typegen!("Snakemake", r###"{json_string}"###, {{
@@ -1153,32 +1141,19 @@ def encode_namedlist(values):
                         .open(path)?;
                     Ok(gag::Redirect::stdout(log)?)
                 }}
-                
-                fn setup_path(&self) -> anyhow::Result<()> {{
-                    use std::env;
-                    if let Some(path) = env::var_os("PATH") {{
-                        let mut paths = env::split_paths(&path).collect::<Vec<_>>();
-                        paths.push(std::path::PathBuf::from("{searchpath}"));
-                        let new_path = env::join_paths(paths)?;
-                        env::set_var("PATH", &new_path);
-                    }}
-                    Ok(())
-                }}
             }}
             
             lazy_static::lazy_static! {{
                 // https://github.com/rust-lang-nursery/lazy-static.rs/issues/153
                 #[allow(non_upper_case_globals)]
                 static ref snakemake: Snakemake = {{
                     let s: Snakemake = serde_json::from_str(r###"{json_string}"###).expect("Failed parsing snakemake JSON");
-                    s.setup_path().expect("Failed setting PATH");
                     s
                 }};
             }}
             // TODO include addendum, if any {{preamble_addendum}}
             """
         ).format(
-            searchpath=searchpath,
             json_string=json_string,
             preamble_addendum=preamble_addendum,
         )