From b4f3191fe2cd62392acc3d6d53a04e384cd57ee9 Mon Sep 17 00:00:00 2001
From: nikpivkin <nikita.pivkin@smartforce.io>
Date: Tue, 17 Dec 2024 14:18:10 +0600
Subject: [PATCH] check for multiple roots

Signed-off-by: nikpivkin <nikita.pivkin@smartforce.io>
---
 pkg/dependency/parser/python/uv/parse.go      | 30 ++++++++++++-------
 pkg/dependency/parser/python/uv/parse_test.go |  5 ++++
 .../python/uv/testdata/uv_multiple_roots.lock | 30 +++++++++++++++++++
 3 files changed, 54 insertions(+), 11 deletions(-)
 create mode 100644 pkg/dependency/parser/python/uv/testdata/uv_multiple_roots.lock

diff --git a/pkg/dependency/parser/python/uv/parse.go b/pkg/dependency/parser/python/uv/parse.go
index ca4925662c66..5ee1ab605bfd 100644
--- a/pkg/dependency/parser/python/uv/parse.go
+++ b/pkg/dependency/parser/python/uv/parse.go
@@ -30,8 +30,7 @@ func (l Lock) directDeps(root Package) map[string]struct{} {
 	return deps
 }
 
-func (l Lock) prodDeps(root Package) map[string]struct{} {
-	packages := l.packages()
+func prodDeps(root Package, packages map[string]Package) map[string]struct{} {
 	visited := make(map[string]struct{})
 	walkPackageDeps(root, packages, visited)
 	return visited
@@ -51,14 +50,25 @@ func walkPackageDeps(pkg Package, packages map[string]Package, visited map[strin
 	}
 }
 
-func (l Lock) root() Package {
+func (l Lock) root() (Package, error) {
+	var pkgs []Package
 	for _, pkg := range l.Packages {
 		if pkg.isRoot() {
-			return pkg
+			pkgs = append(pkgs, pkg)
 		}
 	}
 
-	return Package{}
+	if len(pkgs) > 1 {
+		return Package{}, xerrors.New("uv lockfile contains multiple root projects")
+	}
+
+	// lock file must include root package
+	// cf. https://github.com/astral-sh/uv/blob/f80ddf10b63c3e7b421ca4658e63f97db1e0378c/crates/uv/src/commands/project/lock.rs#L933-L936
+	if len(pkgs) != 1 {
+		return Package{}, xerrors.New("uv lockfile does not contain a root package.")
+	}
+
+	return pkgs[0], nil
 }
 
 type Package struct {
@@ -94,16 +104,14 @@ func (p *Parser) Parse(r xio.ReadSeekerAt) ([]ftypes.Package, []ftypes.Dependenc
 		return nil, nil, xerrors.Errorf("failed to decode uv lock file: %w", err)
 	}
 
-	rootPackage := lock.root()
-	// lock file must include root package
-	// cf. https://github.com/astral-sh/uv/blob/f80ddf10b63c3e7b421ca4658e63f97db1e0378c/crates/uv/src/commands/project/lock.rs#L933-L936
-	if rootPackage.Name == "" {
-		return nil, nil, xerrors.New("uv lockfile does not contain a root package.")
+	rootPackage, err := lock.root()
+	if err != nil {
+		return nil, nil, err
 	}
 
 	packages := lock.packages()
 	directDeps := lock.directDeps(rootPackage)
-	prodDeps := lock.prodDeps(rootPackage)
+	prodDeps := prodDeps(rootPackage, packages)
 
 	var (
 		pkgs ftypes.Packages
diff --git a/pkg/dependency/parser/python/uv/parse_test.go b/pkg/dependency/parser/python/uv/parse_test.go
index 73c02ab3b018..592ce25a6898 100644
--- a/pkg/dependency/parser/python/uv/parse_test.go
+++ b/pkg/dependency/parser/python/uv/parse_test.go
@@ -35,6 +35,11 @@ func TestParser_Parse(t *testing.T) {
 			file:    "testdata/uv_without_root.lock",
 			wantErr: "uv lockfile does not contain a root package",
 		},
+		{
+			name:    "multiple roots",
+			file:    "testdata/uv_multiple_roots.lock",
+			wantErr: "uv lockfile contains multiple root projects",
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/pkg/dependency/parser/python/uv/testdata/uv_multiple_roots.lock b/pkg/dependency/parser/python/uv/testdata/uv_multiple_roots.lock
new file mode 100644
index 000000000000..a0ccda991cb7
--- /dev/null
+++ b/pkg/dependency/parser/python/uv/testdata/uv_multiple_roots.lock
@@ -0,0 +1,30 @@
+version = 1
+requires-python = ">=3.11"
+
+[[package]]
+name = "asyncio"
+version = "3.4.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/da/54/054bafaf2c0fb8473d423743e191fcdf49b2c1fd5e9af3524efbe097bafd/asyncio-3.4.3.tar.gz", hash = "sha256:83360ff8bc97980e4ff25c964c7bd3923d333d177aa4f7fb736b019f26c7cb41", size = 204411 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/22/74/07679c5b9f98a7cb0fc147b1ef1cc1853bc07a4eb9cb5731e24732c5f773/asyncio-3.4.3-py3-none-any.whl", hash = "sha256:c4d18b22701821de07bd6aea8b53d21449ec0ec5680645e5317062ea21817d2d", size = 101767 },
+]
+
+[[package]]
+name = "foo"
+version = "0.1.0"
+source = { virtual = "." }
+dependencies = [
+    { name = "asyncio" },
+]
+
+[package.metadata]
+requires-dist = [{ name = "asyncio", specifier = "==3.4.3" }]
+
+[[package]]
+name = "bar"
+version = "0.1.0"
+source = { virtual = "." }
+dependencies = [
+    { name = "asyncio" },
+]
\ No newline at end of file