fix: flatten pypi packages in the label (#1302)

* fix: flatten pypi packages in the label Signed-off-by: Keming <kemingyang@tensorchord.ai> * delete the fitler for pypi commands Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Dec 14, 2022 · 1cc8c48 · 1cc8c48
1 parent 0a416d1
commit 1cc8c48
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 51 deletions.
diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go
@@ -141,7 +141,11 @@ func (g generalGraph) Labels() (map[string]string, error) {
 		return nil, err
 	}
 	labels[types.ImageLabelAPT] = string(str)
-	str, err = json.Marshal(g.PyPIPackages)
+	pyPackages := []string{}
+	for _, pkg := range g.PyPIPackages {
+		pyPackages = append(pyPackages, pkg...)
+	}
+	str, err = json.Marshal(pyPackages)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
@@ -154,7 +154,11 @@ func (g generalGraph) Labels() (map[string]string, error) {
 		return nil, err
 	}
 	labels[types.ImageLabelAPT] = string(str)
-	str, err = json.Marshal(g.PyPIPackages)
+	pyPackages := []string{}
+	for _, pkg := range g.PyPIPackages {
+		pyPackages = append(pyPackages, pkg...)
+	}
+	str, err = json.Marshal(pyPackages)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/types/envd.go b/pkg/types/envd.go
@@ -17,7 +17,6 @@ package types
 import (
 	"encoding/json"
 	"fmt"
-	"strings"
 
 	"github.com/docker/docker/api/types"
 	"github.com/moby/buildkit/util/system"
@@ -315,20 +314,11 @@ func NewDependencyFromLabels(label map[string]string) (*Dependency, error) {
 		dep.APTPackages = lst
 	}
 	if pypiCommands, ok := label[ImageLabelPyPI]; ok {
-		lst, err := parsePyPICommands(pypiCommands)
+		pkgs, err := parsePyPICommands(pypiCommands)
 		if err != nil {
 			return nil, err
 		}
-		packages := []string{}
-
-		for _, pkgs := range lst {
-			for i, pkg := range pkgs {
-				if !strings.HasPrefix(pkg, "-") && (i == 0 || !strings.HasPrefix(pkgs[i-1], "-")) {
-					packages = append(packages, pkg)
-				}
-			}
-		}
-		dep.PyPIPackages = packages
+		dep.PyPIPackages = pkgs
 	}
 	return &dep, nil
 }
@@ -339,8 +329,8 @@ func parseAPTPackages(lst string) ([]string, error) {
 	return pkgs, err
 }
 
-func parsePyPICommands(lst string) ([][]string, error) {
-	var pkgs [][]string
+func parsePyPICommands(lst string) ([]string, error) {
+	var pkgs []string
 	err := json.Unmarshal([]byte(lst), &pkgs)
 	return pkgs, err
 }

diff --git a/pkg/types/envd_test.go b/pkg/types/envd_test.go