From cd6efc3636047c4e639bbebd4ce823730097672c Mon Sep 17 00:00:00 2001
From: Duc Trung Le <leductrungxf@gmail.com>
Date: Tue, 15 Oct 2024 14:40:01 +0200
Subject: [PATCH] Add package url factory (#108)

* Add package url factory

* Update test
---
 empack/pack.py     | 24 ++++++++++++++++--------
 tests/test_pack.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/empack/pack.py b/empack/pack.py
index 178f496..71a9269 100644
--- a/empack/pack.py
+++ b/empack/pack.py
@@ -5,6 +5,7 @@
 import sys
 from pathlib import Path, PosixPath, PureWindowsPath
 from tempfile import TemporaryDirectory
+from typing import Callable, Optional
 
 from appdirs import user_cache_dir
 
@@ -171,6 +172,7 @@ def pack_env(
     compression_format=ALLOWED_FORMATS[0],
     compresslevel=9,
     outdir=None,
+    package_url_factory: Optional[Callable] = None,
 ):
     with TemporaryDirectory() as tmp_dir:
         #  filter the complete environment
@@ -196,16 +198,22 @@ def pack_env(
             )
 
             base_fname = filename_base_from_meta(pkg_meta)
-            packages_info.append(
-                dict(
-                    name=pkg_meta["name"],
-                    version=pkg_meta["version"],
-                    build=pkg_meta["build"],
-                    filename_stem=base_fname,
-                    filename=f"{base_fname}.tar.{compression_format}",
-                )
+
+            pkg_dict = dict(
+                name=pkg_meta["name"],
+                version=pkg_meta["version"],
+                build=pkg_meta["build"],
+                filename_stem=base_fname,
+                filename=f"{base_fname}.tar.{compression_format}",
             )
 
+            package_url = None
+            if package_url_factory:
+                package_url = package_url_factory(pkg_dict)
+            if package_url is not None:
+                pkg_dict["url"] = package_url
+            packages_info.append(pkg_dict)
+
         # save the list of packages
         env_meta = {
             "prefix": str(relocate_prefix),
diff --git a/tests/test_pack.py b/tests/test_pack.py
index 9035910..b8e39cd 100644
--- a/tests/test_pack.py
+++ b/tests/test_pack.py
@@ -163,6 +163,48 @@ def test_pack_env(tmp_path, packages, relocate_prefix):
             assert pkg_info["name"] == pkg_meta["name"]
 
 
+def test_pack_env_with_url(tmp_path):
+    # create the env at the temporary location
+    packages = ["python=3.11", "numpy"]
+    relocate_prefix = "/"
+    prefix = tmp_path / "env"
+
+    create_environment(
+        prefix=prefix,
+        packages=packages,
+        channels=CHANNELS,
+        relocate_prefix=relocate_prefix,
+        platform="emscripten-wasm32",
+    )
+
+    def url_factory(pkg):
+        return f'http://localhost:1234/{pkg["filename"]}'
+
+    pack_env(
+        env_prefix=prefix,
+        outdir=tmp_path,
+        use_cache=False,
+        compression_format="gz",
+        relocate_prefix=relocate_prefix,
+        file_filters=FILE_FILTERS,
+        compresslevel=1,
+        package_url_factory=url_factory,
+    )
+
+    # check that there is a json with all the packages
+    env_metadata_json_path = tmp_path / "empack_env_meta.json"
+    assert env_metadata_json_path.exists()
+
+    # check that json file contains all packages
+    with open(env_metadata_json_path) as f:
+        env_metadata = json.load(f)
+        packages_metadata = env_metadata["packages"]
+
+    # check that there is a tar.gz file for each package
+    for pkg_info in packages_metadata:
+        assert "http://localhost:1234" in pkg_info["url"]
+
+
 @pytest.mark.parametrize("mount_dir", ["/some", "/some/", "/some/nested", "/"])
 def test_pack_directory(tmp_path, mount_dir):
     # create a directory with some files