Adding support for pip packages

conda_docker/conda.py

@@ -151,7 +151,10 @@ def _precs_from_environment(environment, list_flag, download_dir, user_conda):
                 version=dist.version,
             )
         )
-    return packages
+
+    pip_files = pip_precs_from_environment_prefix(environment)
+    precs_dict = {"precs": packages, "pip_precs": pip_files, "env_path": environment}
+    return precs_dict
 
 
 def precs_from_environment_name(environment, download_dir, user_conda):
@@ -161,6 +164,45 @@ def precs_from_environment_name(environment, download_dir, user_conda):
 def precs_from_environment_prefix(environment, download_dir, user_conda):
     return _precs_from_environment(environment, "--prefix", download_dir, user_conda)
 
+def pip_precs_from_environment_prefix(environment):
+    pip_freeze_output = subprocess.check_output(
+        [f"{environment}/bin/python", "-m", "pip", "freeze"],
+        encoding="utf-8",
+        universal_newlines=True,
+    )
+
+    pip_package_metadatas = []
+    for line in pip_freeze_output.splitlines():
+        line_parts = line.split("==")
+        if len(line_parts) == 2:
+            package_name = line_parts[0]
+            pip_package_metadatas.append(subprocess.check_output(
+                [f"{environment}/bin/python", "-m", "pip", "show", "-f", package_name],
+                encoding="utf-8",
+                universal_newlines=True,
+            ))
+
+    # The output of `pip show -f package` looks something like:
+
+    # Name: foopackage
+    # Location: /opt/conda/envs/some_env/site-packages
+    # Files:
+    #   foopackage/__init__.py 
+    #   foopackage/foo.py
+
+    # We'll use this information to build a list of files to copy into the image.
+
+    pip_files = []
+    for metadata in pip_package_metadatas:
+        pip_file_names = metadata.split("Files:\n")[-1].split('\n')
+        pip_file_names_no_spaces = list(map(lambda path: path.lstrip(), pip_file_names))
+        pip_file_names_no_cache = filter(lambda path: path and not path.endswith(".pyc"), pip_file_names_no_spaces)
+        pip_path_prefix = metadata.split("Location: ")[1].split('\n')[0]
+        pip_path_after_env = pip_path_prefix.split(environment)[1]
+
+        pip_files += list(map(lambda path: os.path.join(pip_path_after_env, path), pip_file_names_no_cache))
+
+    return pip_files
 
 def precs_from_package_specs(
     package_specs,
@@ -285,8 +327,10 @@ def find_solver_conda(solver, user_conda):
 def fetch_precs(download_dir, precs):
     os.makedirs(download_dir, exist_ok=True)
 
+    non_pip_precs = precs["precs"]
+
     records = []
-    for prec in precs:
+    for prec in non_pip_precs:
         package_tarball_full_path = os.path.join(download_dir, prec.fn)
         if package_tarball_full_path.endswith(".tar.bz2"):
             extracted_package_dir = package_tarball_full_path[:-8]
@@ -320,7 +364,8 @@ def fetch_precs(download_dir, precs):
             extracted_package_dir=extracted_package_dir,
         )
         records.append(package_cache_record)
-    return records
+        precs["precs"] = records
+    return precs
 
 
 def write_urls(records, host_pkgs_dir, channels_remap):
@@ -647,6 +692,17 @@ def build_docker_environment(
     with timer(LOGGER, "writing docker file"):
         image.write_filename(output_filename)
 
+def copy_pip_packages(pip_targ_dir, pip_env, pip_files):
+    for pip_file in pip_files:
+        # Preserve the original pip file paths from the download dir to the target dir
+        # by concatenating the start of the target dir with the end of source file path
+        # (essentially just switching the source prefix with the target prefix)
+        pip_path_source = os.path.join(pip_env, pip_file[1:]) # /opt/conda/envs/test/path_to_file/file_name
+        if os.path.isfile(pip_path_source):
+            pip_path_target = os.path.join(pip_targ_dir, pip_file[1:pip_file.rindex('/')]) # /tmp/tmpzlujmuh3/opt/conda/path_to_file
+            if not os.path.isdir(pip_path_target):
+                os.makedirs(pip_path_target)
+            shutil.copy(pip_path_source, pip_path_target)
 
 def build_docker_environment_image(
     base_image: Image,
@@ -662,25 +718,42 @@ def build_docker_environment_image(
     base_image.name = output_image_name
     base_image.tag = output_image_tag
 
+    non_pip_packages = records["precs"]
+    pip_file_paths = records["pip_precs"]
+    env_path = records["env_path"]
+
     with tempfile.TemporaryDirectory() as tmpdir:
         LOGGER.info("building conda environment")
+
         with timer(LOGGER, "building conda environment"):
             chroot_install(
                 str(tmpdir),
-                records,
+                non_pip_packages,
                 default_prefix,
                 download_dir,
                 user_conda,
                 channels_remap,
             )
 
+            pip_targ_dir = os.path.join(str(tmpdir), "opt", "conda")
+
+            copy_pip_packages(
+                pip_targ_dir,
+                env_path,
+                pip_file_paths,
+            )
+
         add_conda_layers(
             base_image,
             str(tmpdir),
             arcpath="/",
             filter=conda_file_filter(),
-            records=records,
+            records=non_pip_packages,
             layering_strategy=layering_strategy,
         )
+        add_single_conda_layer(
+            base_image,
+            pip_targ_dir,
+        )
 
         return base_image