refactor(python): use set

Signed-off-by: knqyf263 <[email protected]>
aquasecurity · Dec 23, 2024 · 75667ef · 75667ef
1 parent f6176b4
commit 75667ef
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 19 deletions.
diff --git a/pkg/dependency/parser/python/pyproject/pyproject.go b/pkg/dependency/parser/python/pyproject/pyproject.go
@@ -8,6 +8,7 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/aquasecurity/trivy/pkg/dependency/parser/python"
+	"github.com/aquasecurity/trivy/pkg/set"
 )
 
 type PyProject struct {
@@ -19,25 +20,27 @@ type Tool struct {
 }
 
 type Poetry struct {
-	Dependencies dependencies     `toml:"dependencies"`
+	Dependencies Dependencies     `toml:"dependencies"`
 	Groups       map[string]Group `toml:"group"`
 }
 
 type Group struct {
-	Dependencies dependencies `toml:"dependencies"`
+	Dependencies Dependencies `toml:"dependencies"`
 }
 
-type dependencies map[string]struct{}
+type Dependencies struct {
+	set.Set[string]
+}
 
-func (d *dependencies) UnmarshalTOML(data any) error {
+func (d *Dependencies) UnmarshalTOML(data any) error {
 	m, ok := data.(map[string]any)
 	if !ok {
 		return xerrors.Errorf("dependencies must be map, but got: %T", data)
 	}
 
-	*d = lo.MapEntries(m, func(pkgName string, _ any) (string, struct{}) {
-		return python.NormalizePkgName(pkgName), struct{}{}
-	})
+	d.Set = set.New[string](lo.MapToSlice(m, func(pkgName string, _ any) string {
+		return python.NormalizePkgName(pkgName)
+	})...)
 	return nil
 }
 

diff --git a/pkg/dependency/parser/python/pyproject/pyproject_test.go b/pkg/dependency/parser/python/pyproject/pyproject_test.go
@@ -2,6 +2,7 @@ package pyproject_test
 
 import (
 	"fmt"
+	"github.com/aquasecurity/trivy/pkg/set"
 	"os"
 	"testing"
 
@@ -24,21 +25,18 @@ func TestParser_Parse(t *testing.T) {
 			want: pyproject.PyProject{
 				Tool: pyproject.Tool{
 					Poetry: pyproject.Poetry{
-						Dependencies: map[string]struct{}{
-							"flask":      {},
-							"python":     {},
-							"requests":   {},
-							"virtualenv": {},
+						Dependencies: pyproject.Dependencies{
+							Set: set.New[string]("flask", "python", "requests", "virtualenv"),
 						},
 						Groups: map[string]pyproject.Group{
 							"dev": {
-								Dependencies: map[string]struct{}{
-									"pytest": {},
+								Dependencies: pyproject.Dependencies{
+									Set: set.New[string]("pytest"),
 								},
 							},
 							"lint": {
-								Dependencies: map[string]struct{}{
-									"ruff": {},
+								Dependencies: pyproject.Dependencies{
+									Set: set.New[string]("ruff"),
 								},
 							},
 						},

diff --git a/pkg/fanal/analyzer/language/python/poetry/poetry.go b/pkg/fanal/analyzer/language/python/poetry/poetry.go
@@ -106,7 +106,7 @@ func (a poetryAnalyzer) mergePyProject(fsys fs.FS, dir string, app *types.Applic
 
 	// Identify the direct/transitive dependencies
 	for i, pkg := range app.Packages {
-		if _, ok := project.Tool.Poetry.Dependencies[pkg.Name]; ok {
+		if project.Tool.Poetry.Dependencies.Contains(pkg.Name) {
 			app.Packages[i].Relationship = types.RelationshipDirect
 		} else {
 			app.Packages[i].Indirect = true
@@ -130,11 +130,11 @@ func filterProdPackages(project pyproject.PyProject, app *types.Application) {
 		if group == "dev" {
 			continue
 		}
-		deps = lo.Assign(deps, groupDeps.Dependencies)
+		deps.Set = deps.Union(groupDeps.Dependencies)
 	}
 
 	for _, pkg := range packages {
-		if _, prodDep := deps[pkg.Name]; !prodDep {
+		if !deps.Contains(pkg.Name) {
 			continue
 		}
 		walkPackageDeps(pkg.ID, packages, visited)