feat(python): add support for uv optional dependencies

Signed-off-by: nikpivkin <[email protected]>
aquasecurity · Dec 19, 2024 · 3ce2c3c · 3ce2c3c
1 parent 6800ba5
commit 3ce2c3c
Show file tree

Hide file tree

Showing 5 changed files with 365 additions and 66 deletions.
diff --git a/pkg/dependency/parser/python/uv/parse.go b/pkg/dependency/parser/python/uv/parse.go
@@ -22,19 +22,6 @@ func (l Lock) packages() map[string]Package {
 	})
 }
 
-func (p Package) directDeps() map[string]struct{} {
-	deps := make(map[string]struct{})
-	for _, dep := range p.Dependencies {
-		deps[dep.Name] = struct{}{}
-	}
-	for _, groupDeps := range p.DevDependencies {
-		for _, dep := range groupDeps {
-			deps[dep.Name] = struct{}{}
-		}
-	}
-	return deps
-}
-
 func prodDeps(root Package, packages map[string]Package) map[string]struct{} {
 	visited := make(map[string]struct{})
 	walkPackageDeps(root, packages, visited)
@@ -46,8 +33,8 @@ func walkPackageDeps(pkg Package, packages map[string]Package, visited map[strin
 		return
 	}
 	visited[pkg.Name] = struct{}{}
-	for _, dep := range pkg.Dependencies {
-		depPkg, exists := packages[dep.Name]
+	for depName := range pkg.nonDevDeps() {
+		depPkg, exists := packages[depName]
 		if !exists {
 			continue
 		}
@@ -73,11 +60,38 @@ func (l Lock) root() (Package, error) {
 }
 
 type Package struct {
-	Name            string                  `toml:"name"`
-	Version         string                  `toml:"version"`
-	Source          Source                  `toml:"source"`
-	Dependencies    []Dependency            `toml:"dependencies"`
-	DevDependencies map[string][]Dependency `toml:"dev-dependencies"`
+	Name                 string                  `toml:"name"`
+	Version              string                  `toml:"version"`
+	Source               Source                  `toml:"source"`
+	Dependencies         Dependencies            `toml:"dependencies"`
+	DevDependencies      map[string]Dependencies `toml:"dev-dependencies"`
+	OptionalDependencies map[string]Dependencies `toml:"optional-dependencies"`
+}
+
+func (p Package) directDeps() map[string]struct{} {
+	deps := p.nonDevDeps()
+	for _, groupDeps := range p.DevDependencies {
+		deps = lo.Assign(deps, groupDeps.toSet())
+	}
+	return deps
+}
+
+func (p Package) nonDevDeps() map[string]struct{} {
+	deps := p.Dependencies.toSet()
+	for _, groupDeps := range p.OptionalDependencies {
+		deps = lo.Assign(deps, groupDeps.toSet())
+	}
+	return deps
+}
+
+type Dependencies []Dependency
+
+func (d Dependencies) toSet() map[string]struct{} {
+	deps := make(map[string]struct{})
+	for _, dep := range d {
+		deps[dep.Name] = struct{}{}
+	}
+	return deps
 }
 
 // https://github.com/astral-sh/uv/blob/f7d647e81d7e1e3be189324b06024ed2057168e6/crates/uv-resolver/src/lock/mod.rs#L572-L579

diff --git a/pkg/dependency/parser/python/uv/parse_testcase.go b/pkg/dependency/parser/python/uv/parse_testcase.go
@@ -9,30 +9,43 @@ var (
 	// uv init normal && cd normal
 	// uv add requests==2.32.0
 	// uv add --group dev pytest==8.3.4
+	// uv add httpx==0.28.1 --extra socks
+	// uv add orjson==3.10.12 --optional json
 	// apk add jq
 	// uv pip list --format json |jq -c 'sort_by(.name) | .[] | {"ID": (.name + "@" + .version), "Name": .name, "Version": .version}' | sed 's/$/,/' | sed 's/\"\([^"]*\)\":/\1:/g'
 
 	// add a root project
 	// fill in the relationships between the packages
 	uvNormal = []ftypes.Package{
 		{ID: "[email protected]", Name: "normal", Version: "0.1.0", Relationship: ftypes.RelationshipRoot},
+		{ID: "[email protected]", Name: "httpx", Version: "0.28.1", Relationship: ftypes.RelationshipDirect},
+		{ID: "[email protected]", Name: "orjson", Version: "3.10.12", Relationship: ftypes.RelationshipDirect},
 		{ID: "[email protected]", Name: "pytest", Version: "8.3.4", Relationship: ftypes.RelationshipDirect, Dev: true},
 		{ID: "[email protected]", Name: "requests", Version: "2.32.0", Relationship: ftypes.RelationshipDirect},
+		{ID: "[email protected]", Name: "anyio", Version: "4.7.0", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "certifi", Version: "2024.12.14", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "charset-normalizer", Version: "3.4.0", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "colorama", Version: "0.4.6", Relationship: ftypes.RelationshipIndirect, Dev: true},
-		{ID: "[email protected]", Name: "exceptiongroup", Version: "1.2.2", Relationship: ftypes.RelationshipIndirect, Dev: true},
+		{ID: "[email protected]", Name: "exceptiongroup", Version: "1.2.2", Relationship: ftypes.RelationshipIndirect},
+		{ID: "[email protected]", Name: "h11", Version: "0.14.0", Relationship: ftypes.RelationshipIndirect},
+		{ID: "[email protected]", Name: "httpcore", Version: "1.0.7", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "idna", Version: "3.10", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "iniconfig", Version: "2.0.0", Relationship: ftypes.RelationshipIndirect, Dev: true},
 		{ID: "[email protected]", Name: "packaging", Version: "24.2", Relationship: ftypes.RelationshipIndirect, Dev: true},
 		{ID: "[email protected]", Name: "pluggy", Version: "1.5.0", Relationship: ftypes.RelationshipIndirect, Dev: true},
+		{ID: "[email protected]", Name: "sniffio", Version: "1.3.1", Relationship: ftypes.RelationshipIndirect},
+		{ID: "[email protected]", Name: "socksio", Version: "1.0.0", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "tomli", Version: "2.2.1", Relationship: ftypes.RelationshipIndirect, Dev: true},
+		{ID: "[email protected]", Name: "typing-extensions", Version: "4.12.2", Relationship: ftypes.RelationshipIndirect},
 		{ID: "[email protected]", Name: "urllib3", Version: "2.2.3", Relationship: ftypes.RelationshipIndirect},
 	}
 
 	// add a root project
 	uvNormalDeps = []ftypes.Dependency{
-		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]"}},
+		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]", "[email protected]", "[email protected]"}},
+		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]"}},
+		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]", "[email protected]", "[email protected]", "[email protected]"}},
+		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]", "[email protected]", "[email protected]"}},
 		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]", "[email protected]", "[email protected]", "[email protected]", "[email protected]"}},
 		{ID: "[email protected]", DependsOn: []string{"[email protected]", "[email protected]", "[email protected]", "[email protected]"}},
 	}