bug: fix pypi package information (#959)

* fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * add test Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
tensorchord · Oct 9, 2022 · 91ea50c · 91ea50c
1 parent 5e8825c
commit 91ea50c
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 5 deletions.
diff --git a/pkg/types/envd.go b/pkg/types/envd.go
@@ -17,6 +17,7 @@ package types
 import (
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/docker/docker/api/types"
 	"github.com/moby/buildkit/util/system"
@@ -249,12 +250,19 @@ func newDependencyFromLabels(label map[string]string) (*Dependency, error) {
 		}
 		dep.APTPackages = lst
 	}
-	if pkgs, ok := label[ImageLabelPyPI]; ok {
-		lst, err := parsePyPIPackages(pkgs)
+	if pypiCommands, ok := label[ImageLabelPyPI]; ok {
+		lst, err := parsePyPICommands(pypiCommands)
 		if err != nil {
 			return nil, err
 		}
-		dep.PyPIPackages = lst
+		packages := []string{}
+
+		for i, pkg := range lst {
+			if !strings.HasPrefix(pkg, "-") && (i == 0 || !strings.HasPrefix(lst[i-1], "-")) {
+				packages = append(packages, pkg)
+			}
+		}
+		dep.PyPIPackages = packages
 	}
 	return &dep, nil
 }
@@ -265,7 +273,7 @@ func parseAPTPackages(lst string) ([]string, error) {
 	return pkgs, err
 }
 
-func parsePyPIPackages(lst string) ([]string, error) {
+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
@@ -0,0 +1,35 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package types
+
+import (
+	"encoding/json"
+
+	g "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = g.Describe("pypi label", func() {
+	g.When("pypi command contains slashes", func() {
+		g.It("should ignore such commands", func() {
+			labels := make(map[string]string)
+			s, _ := json.Marshal([]string{"numpy", "torch", "--extra-url", "https://download.pytorch.org/whl/torch_stable.html"})
+			labels[ImageLabelPyPI] = string(s)
+			d, err := newDependencyFromLabels(labels)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(d.PyPIPackages).To(Equal([]string{"numpy", "torch"}))
+		})
+	})
+})
diff --git a/pkg/types/label.go b/pkg/types/label.go
@@ -23,7 +23,7 @@ const (
 	ImageLabelVendor    = "ai.tensorchord.envd.vendor"
 	ImageLabelGPU       = "ai.tensorchord.envd.gpu"
 	ImageLabelAPT       = "ai.tensorchord.envd.apt.packages"
-	ImageLabelPyPI      = "ai.tensorchord.envd.pypi.packages"
+	ImageLabelPyPI      = "ai.tensorchord.envd.pypi.commands"
 	ImageLabelR         = "ai.tensorchord.envd.r.packages"
 	ImageLabelCUDA      = "ai.tensorchord.envd.gpu.cuda"
 	ImageLabelCUDNN     = "ai.tensorchord.envd.gpu.cudnn"

diff --git a/pkg/types/types_suite_test.go b/pkg/types/types_suite_test.go
@@ -0,0 +1,27 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package types
+
+import (
+	"testing"
+
+	g "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestLabelTypes(t *testing.T) {
+	RegisterFailHandler(g.Fail)
+	g.RunSpecs(t, "label types Suite")
+}