From afd4d0a637a64fa17a2e385c88185283e71f2631 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Mon, 14 Oct 2024 16:04:43 +0100
Subject: [PATCH] Support extracting minimum constraints for monorepo

---
 .../base-setup/create_constraints_file.py     | 51 +++++++++++++++----
 .../actions/base-setup/setup_constraints.sh   |  4 +-
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/.github/actions/base-setup/create_constraints_file.py b/.github/actions/base-setup/create_constraints_file.py
index 2d26770..b7a384b 100644
--- a/.github/actions/base-setup/create_constraints_file.py
+++ b/.github/actions/base-setup/create_constraints_file.py
@@ -1,21 +1,52 @@
 import sys
+from typing import cast, List
 from pathlib import Path
-from zipfile import ZipFile
 
+from build.util import project_wheel_metadata
 from packaging.requirements import Requirement
+from importlib.metadata import PackageMetadata
 
 output_file = sys.argv[-2]
-fname = sys.argv[-1]
+top_level_project_dir = sys.argv[-1]
 constraints = {}
 
-archive = ZipFile(fname)
-reqs = []
-for f in archive.namelist():
-    if f.endswith("METADATA"):
-        for li in archive.open(f).read().decode("utf-8").split("\n"):
-            if li.startswith("Requires-Dist"):
-                reqs.append(li.replace("Requires-Dist: ", ""))
-archive.close()
+
+def extract_dependencies(project_dir):
+    reqs = []
+    print(f'Extracting metadata from wheel for {project_dir}...')
+    metadata = project_wheel_metadata(source_dir=project_dir)
+    reqs.extend(get_requires_dist(metadata))
+
+    # extract requirements from local dependencies specified with file: protocol
+    # to support the mono-repo usecase
+    local_projects = {}
+    for req in reqs:
+        r = Requirement(req)
+        if r.url and r.url.startswith('file://'):
+            path = r.url.replace('file://', '')
+            local_projects[r.name] = path
+
+    reqs_from_local_dependencies = []
+    for dependency_name, path in local_projects.items():
+        print(f'Discovering constraints in local {r.name} package under {path}')
+        sub_dependencies = extract_dependencies(path)
+        # filter out dependencies between local packages (e.g. jupyter-ai depends on
+        # a fixed minimum version of `jupyter-ai-magics`, but here we want to test
+        # the latest version against its third-party dependencies - not the old one).
+        sub_dependencies = [
+            req
+            for req in sub_dependencies
+            if Requirement(req).name not in local_projects
+        ]
+        reqs_from_local_dependencies.extend(sub_dependencies)
+    return reqs + reqs_from_local_dependencies
+
+
+def get_requires_dist(metadata: PackageMetadata) -> List[str]:
+    return cast(List[str], metadata.get_all('Requires-Dist'))
+
+
+reqs = extract_dependencies(top_level_project_dir)
 
 for req in reqs:
     r = Requirement(req)
diff --git a/.github/actions/base-setup/setup_constraints.sh b/.github/actions/base-setup/setup_constraints.sh
index 192be2c..e09bdc6 100644
--- a/.github/actions/base-setup/setup_constraints.sh
+++ b/.github/actions/base-setup/setup_constraints.sh
@@ -2,10 +2,8 @@
 python -m venv $HOME/.venv
 source $HOME/.venv/bin/activate
 pip install build packaging pkginfo
-mkdir $HOME/dist
-python -m build --outdir $HOME/dist --wheel .
 
 SCRIPT_DIR=$(cd $(dirname "${BASH_SOURCE[0]}") && pwd)
-python $SCRIPT_DIR/create_constraints_file.py $HOME/constraints.txt $HOME/dist/*.whl
+python $SCRIPT_DIR/create_constraints_file.py $HOME/constraints.txt $(pwd)
 cat $HOME/constraints.txt
 echo "PIP_CONSTRAINT=$HOME/constraints.txt" >> $GITHUB_ENV