Merge pull request #13 from mariusvniekerk/pkgs_dir_workaround

Allow conda-lock to work when pkgs_dir already contains artifacts

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+ignore = E203, E266, E501, W503, F403, F401
+max-line-length = 89
+max-complexity = 18
+select = B,C,E,F,W,T4,B9

.gitignore

@@ -1,7 +1,10 @@
 *.egg-info
 *.eggs
 *.pyc
+*.tar.gz
 .mypy_cache/
 .vscode/
+.idea/
 build/
 dist/
+

.pre-commit-config.yaml

@@ -1,4 +1,15 @@
 repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v2.3.0
+  hooks:
+    - id: trailing-whitespace
+    - id: check-ast
+
+-   repo: https://gitlab.com/pycqa/flake8
+    rev: 3.7.1
+    hooks:
+    -   id: flake8
+
 - repo: https://github.com/pre-commit/mirrors-isort
   rev: v4.3.21
   hooks:

.travis.yml

@@ -65,7 +65,7 @@ script:
       pushd docs
       make clean html linkcheck
       popd
-      if [[ -z "$TRAVIS_TAG" ]]; then 
+      if [[ -z "$TRAVIS_TAG" ]]; then
         python -m doctr deploy --build-tags --key-path github_deploy_key.enc --built-docs docs/_build/html dev
       else
         python -m doctr deploy --build-tags --key-path github_deploy_key.enc --built-docs docs/_build/html "version-$TRAVIS_TAG"

Dockerfile

@@ -0,0 +1,5 @@
+FROM continuumio/miniconda:latest
+
+RUN pip install conda-lock
+
+ENTRYPOINT conda-lock

README.md

@@ -17,7 +17,7 @@ various places
 
 ### Dockerfile example
 
-In order to use conda-lock in a docker-style context you want to add the lockfile to the 
+In order to use conda-lock in a docker-style context you want to add the lockfile to the
 docker container.  In order to refresh the lock file just run `conda-lock` again.
 ```
   Dockerfile

conda_lock/conda_lock.py

@@ -2,19 +2,27 @@
 Somewhat hacky solution to create conda lock files.
 """
 
+import atexit
 import json
+import logging
 import os
 import pathlib
+import platform
+import shutil
+import stat
 import subprocess
 import sys
 import tempfile
 
-from typing import Dict, List
+from typing import Dict, List, MutableSequence, Optional, Set, Tuple, Union
 
 import requests
 import yaml
 
 
+PathLike = Union[str, pathlib.Path]
+
+
 if not (sys.version_info.major >= 3 and sys.version_info.minor >= 6):
     print("conda_lock needs to run under python >=3.6")
     sys.exit(1)
@@ -23,42 +31,130 @@
 DEFAULT_PLATFORMS = ["osx-64", "linux-64", "win-64"]
 
 
-def solve_specs_for_arch(channels: List[str], specs: List[str], platform: str) -> str:
+def ensure_conda(conda_executable: Optional[str]):
+    if conda_executable:
+        if pathlib.Path(conda_executable).exists():
+            return conda_executable
+
+    conda_executable = shutil.which("conda")
+    if conda_executable:
+        return conda_executable
+
+    conda_executable = shutil.which("conda.exe")
+    if conda_executable:
+        return conda_executable
+
+    logging.info(
+        "No existing conda installation found.  Installing the standalone conda solver"
+    )
+    return install_conda_exe()
+
+
+def install_conda_exe():
+    conda_exe_prefix = "https://repo.anaconda.com/pkgs/misc/conda-execs"
+    if platform.system() == "Linux":
+        conda_exe_file = "conda-latest-linux-64.exe"
+    elif platform.system() == "Darwin":
+        conda_exe_file = "conda-latest-osx-64.exe"
+    elif platform.system() == "NT":
+        conda_exe_file = "conda-latest-win-64.exe"
+    else:
+        # TODO: Support windows here
+        raise ValueError(f"Unsupported platform: {platform.system()}")
+
+    resp = requests.get(f"{conda_exe_prefix}/{conda_exe_file}", allow_redirects=True)
+    resp.raise_for_status()
+    target_filename = os.path.expanduser(pathlib.Path(__file__).parent / "conda.exe")
+    with open(target_filename, "wb") as fo:
+        fo.write(resp.content)
+    st = os.stat(target_filename)
+    os.chmod(target_filename, st.st_mode | stat.S_IXUSR)
+    return target_filename
+
+
+CONDA_PKGS_DIRS = None
+
+
+def conda_pkgs_dir():
+    global CONDA_PKGS_DIRS
+    if CONDA_PKGS_DIRS is None:
+        CONDA_PKGS_DIRS = tempfile.TemporaryDirectory()
+        atexit.register(CONDA_PKGS_DIRS.cleanup)
+        return CONDA_PKGS_DIRS
+
+
+def conda_env_override():
     env = dict(os.environ)
-    with tempfile.TemporaryDirectory() as CONDA_PKGS_DIRS:
-        env.update(
-            {
-                "CONDA_SUBDIR": platform,
-                "CONDA_PKGS_DIRS": CONDA_PKGS_DIRS,
-                "CONDA_UNSATISFIABLE_HINTS_CHECK_DEPTH": "0",
-                "CONDA_ADD_PIP_AS_PYTHON_DEPENDENCY": "False",
-            }
+    env.update(
+        {
+            "CONDA_SUBDIR": platform,
+            "CONDA_PKGS_DIRS": conda_pkgs_dir(),
+            "CONDA_UNSATISFIABLE_HINTS_CHECK_DEPTH": "0",
+            "CONDA_ADD_PIP_AS_PYTHON_DEPENDENCY": "False",
+        }
+    )
+    return env
+
+
+def solve_specs_for_arch(
+    conda: PathLike, channels: List[str], specs: List[str], platform: str
+) -> dict:
+    args: MutableSequence[PathLike] = [
+        conda,
+        "create",
+        "--prefix",
+        pathlib.Path(conda_pkgs_dir()).joinpath("prefix"),
+        "--override-channels",
+        "--dry-run",
+        "--json",
+    ]
+    for channel in channels:
+        args.extend(["--channel", channel])
+    args.extend(specs)
+
+    try:
+        proc = subprocess.run(
+            args, env=conda_env_override(), capture_output=True, encoding="utf8"
         )
+        proc.check_returncode()
+    except subprocess.CalledProcessError:
+        err_json = json.loads(proc.stdout)
+        print(err_json["message"])
+        print("\n")
+        print(f"Could not lock the environment for platform {platform}")
+        sys.exit(1)
 
-        args = [
-            "conda",
-            "create",
-            "--prefix",
-            str(pathlib.Path(CONDA_PKGS_DIRS).joinpath("prefix")),
-            "--override-channels",
-            "--dry-run",
-            "--json",
-        ]
-        for channel in channels:
-            args.extend(["--channel", channel])
-        args.extend(specs)
+    return json.loads(proc.stdout)
+
+
+def search_for_md5s(conda: PathLike, package_specs: List[dict]):
+    """Use conda-search to determine the md5 metadata that we need.
 
-        try:
-            proc = subprocess.run(args, env=env, capture_output=True, encoding="utf8")
-            proc.check_returncode()
-        except subprocess.CalledProcessError:
-            err_json = json.loads(proc.stdout)
-            print(err_json["message"])
-            print("\n")
-            print(f"Could not lock the environment for platform {platform}")
-            sys.exit(1)
+    This is only needed if pkgs_dirs is set in condarc.
+    Sadly this is going to be slow since we need to fetch each result individually
+    due to the cli of conda search
 
-        return json.loads(proc.stdout)
+    """
+    found: Set[str] = set()
+    packages: List[Tuple[str, str]] = [
+        *[(d["name"], f"{d['name']}[url={d['url_conda']}]") for d in package_specs],
+        *[(d["name"], f"{d['name']}[url={d['url']}]") for d in package_specs],
+    ]
+
+    for name, spec in packages:
+        if name in found:
+            continue
+        out = subprocess.run(
+            ["conda", "search", "--use-index-cache", "--json", spec],
+            encoding="utf8",
+            capture_output=True,
+            env=conda_env_override(),
+        )
+        content = json.loads(out.stdout)
+        if name in content:
+            assert len(content[name]) == 1
+            yield content[name][0]
+            found.add(name)
 
 
 def parse_environment_file(environment_file: pathlib.Path) -> Dict:
@@ -83,26 +179,50 @@ def fn_to_dist_name(fn: str) -> str:
     return fn
 
 
-def make_lock_files(platforms, channels, specs):
-    for platform in platforms:
-        print(f"generating lockfile for {platform}", file=sys.stderr)
+def make_lock_files(
+    conda: PathLike, platforms: List[str], channels: List[str], specs: List[str]
+):
+    for plat in platforms:
+        print(f"generating lockfile for {plat}", file=sys.stderr)
         dry_run_install = solve_specs_for_arch(
-            platform=platform, channels=channels, specs=specs
+            conda=conda, platform=plat, channels=channels, specs=specs
         )
-        with open(f"conda-{platform}.lock", "w") as fo:
-            fo.write(f"# platform: {platform}\n")
+        with open(f"conda-{plat}.lock", "w") as fo:
+            fo.write(f"# platform: {plat}\n")
             fo.write("@EXPLICIT\n")
-            urls = {
-                fn_to_dist_name(pkg["fn"]): pkg["url"]
-                for pkg in dry_run_install["actions"]["FETCH"]
-            }
-            md5s = {
-                fn_to_dist_name(pkg["fn"]): pkg["md5"]
-                for pkg in dry_run_install["actions"]["FETCH"]
+            link_actions = dry_run_install["actions"]["LINK"]
+            for link in link_actions:
+                link[
+                    "url_base"
+                ] = f"{link['base_url']}/{link['platform']}/{link['dist_name']}"
+                link["url"] = f"{link['url_base']}.tar.bz2"
+                link["url_conda"] = f"{link['url_base']}.conda"
+            link_dists = {link["dist_name"] for link in link_actions}
+
+            fetch_actions = dry_run_install["actions"]["FETCH"]
+            import pprint
+
+            pprint.pprint(("FETCH", fetch_actions))
+
+            fetch_by_dist_name = {
+                fn_to_dist_name(pkg["fn"]): pkg for pkg in fetch_actions
             }
-            for pkg in dry_run_install["actions"]["LINK"]:
-                url = urls[pkg["dist_name"]]
-                md5 = md5s[pkg["dist_name"]]
+
+            non_fetch_packages = link_dists - set(fetch_by_dist_name)
+            if len(non_fetch_packages) > 0:
+                print(link_dists - set(fetch_by_dist_name))
+                for search_res in search_for_md5s(
+                    conda,
+                    [l for l in link_actions if l["dist_name"] in non_fetch_packages],
+                ):
+                    print(search_res)
+                    dist_name = fn_to_dist_name(search_res["fn"])
+                    print(dist_name)
+                    fetch_by_dist_name[dist_name] = search_res
+
+            for pkg in link_actions:
+                url = fetch_by_dist_name[pkg["dist_name"]]["url"]
+                md5 = fetch_by_dist_name[pkg["dist_name"]]["md5"]
                 r = requests.head(url, allow_redirects=True)
                 url = r.url
                 fo.write(f"{url}#{md5}")
@@ -116,10 +236,36 @@ def make_lock_files(platforms, channels, specs):
     print("", file=sys.stderr)
 
 
+def main_on_docker(env_file, platforms):
+    env_path = pathlib.Path(env_file)
+    platform_arg = []
+    for p in platforms:
+        platform_arg.extend(["--platform", p])
+
+    subprocess.check_output(
+        [
+            "docker",
+            "run",
+            "--rm",
+            "-v",
+            f"{str(env_path.parent)}:/work:rwZ",
+            "--workdir",
+            "/work",
+            "conda-lock:latest",
+            "--file",
+            env_path.name,
+            *platform_arg,
+        ]
+    )
+
+
 def main():
     import argparse
 
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c", "--conda", default=None, help="path to the conda executable to use."
+    )
     parser.add_argument(
         "-p",
         "--platform",
@@ -133,12 +279,25 @@ def main():
         default="environment.yml",
         help="path to a conda environment specification",
     )
+    parser.add_argument(
+        "-m",
+        "--mode",
+        choices=["default", "docker"],
+        default="default",
+        help="""
+            Run this conda-lock in an isolated docker container.  This may be
+            required to account for some issues where conda-lock condflicts with
+            existing condarc configurations.
+            """,
+    )
 
     args = parser.parse_args()
 
     environment_file = pathlib.Path(args.file)
     desired_env = parse_environment_file(environment_file)
+    conda_exe = ensure_conda(args.conda)
     make_lock_files(
+        conda=conda_exe,
         channels=desired_env["channels"] or [],
         specs=desired_env["specs"],
         platforms=args.platform or DEFAULT_PLATFORMS,

docs/source/conf.py

@@ -58,7 +58,7 @@
 # |version| and |release|, also used in various other places throughout the
 # built documents.
 #
-import conda_lock
+import conda_lock  # noqa: E402
 
 
 version = release = conda_lock.__version__

docs/source/index.rst

@@ -29,6 +29,6 @@ Generate the lockfiles,
 
 Create an environment from the lockfile
 
-  .. code-block:: bash     
+  .. code-block:: bash
 
      conda create --name my-locked-env --file conda-linux-64.lock