vpython: Add wheel parsing/management package.

vpython/wheel/wheel.go

+// Copyright 2017 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package wheel
+
+import (
+        "fmt"
+        "os"
+        "path/filepath"
+        "strings"
+
+        "github.com/luci/luci-go/common/errors"
+)
+
+// Name is a parsed Python wheel name, defined here:
+// https://www.python.org/dev/peps/pep-0427/#file-name-convention
+//
+// {distribution}-{version}(-{build tag})?-{python tag}-{abi tag}-\
+// {platform tag}.whl .
+type Name struct {
+        Distribution string
+        Version      string
+        BuildTag     string
+        PythonTag    string
+        ABITag       string
+        PlatformTag  string
+}
+
+func (wn *Name) String() string {
+        parts := make([]string, 0, 6)
+        parts = append(parts, []string{
+                wn.Distribution,
+                wn.Version,
+        }...)
+        if wn.BuildTag != "" {
+                parts = append(parts, wn.BuildTag)
+        }
+        parts = append(parts, []string{
+                wn.PythonTag,
+                wn.ABITag,
+                wn.PlatformTag,
+        }...)
+        return strings.Join(parts, "-") + ".whl"
+}
+
+// ParseName parses a wheel Name from its filename.
+func ParseName(v string) (wn Name, err error) {
+        base := strings.TrimSuffix(v, ".whl")
+        if len(base) == len(v) {
+                err = errors.Reason("missing .whl suffix").Err()
+                return
+        }
+
+        skip := 0
+        switch parts := strings.Split(base, "-"); len(parts) {
+        case 6:
+                // Extra part: build tag.
+                wn.BuildTag = parts[2]
+                skip = 1
+                fallthrough
+
+        case 5:
+                wn.Distribution = parts[0]
+                wn.Version = parts[1]
+                wn.PythonTag = parts[2+skip]
+                wn.ABITag = parts[3+skip]
+                wn.PlatformTag = parts[4+skip]
+
+        default:
+                err = errors.Reason("unknown number of segments (%(segments)d)").
+                        D("segments", len(parts)).
+                        Err()
+                return
+        }
+        return
+}
+
+// ScanDir identifies all wheel files in the immediate directory dir and
+// returns their parsed wheel names.
+func ScanDir(dir string) ([]Name, error) {
+        globPattern := filepath.Join(dir, "*.whl")
+        matches, err := filepath.Glob(globPattern)
+        if err != nil {
+                return nil, errors.Annotate(err).Reason("failed to list wheel directory: %(dir)s").
+                        D("dir", dir).
+                        D("pattern", globPattern).
+                        Err()
+        }
+
+        names := make([]Name, 0, len(matches))
+        for _, match := range matches {
+                switch st, err := os.Stat(match); {
+                case err != nil:
+                        return nil, errors.Annotate(err).Reason("failed to stat wheel: %(path)s").
+                                D("path", match).
+                                Err()
+
+                case st.IsDir():
+                        // Ignore directories.
+                        continue
+
+                default:
+                        // A ".whl" file.
+                        name := filepath.Base(match)
+                        wheelName, err := ParseName(name)
+                        if err != nil {
+                                return nil, errors.Annotate(err).Reason("failed to parse wheel from: %(name)s").
+                                        D("name", name).
+                                        D("dir", dir).
+                                        Err()
+                        }
+                        names = append(names, wheelName)
+                }
+        }
+        return names, nil
+}
+
+// WriteRequirementsFile writes a valid "requirements.txt"-style pip
+// requirements file containing the supplied wheels.
+//
+// The generated requirements will request the exact wheel senver version (using
+// "==").
+func WriteRequirementsFile(path string, wheels []Name) (err error) {
+        fd, err := os.Create(path)
+        if err != nil {
+                return errors.Annotate(err).Reason("failed to create requirements file").Err()
+        }
+        defer func() {
+                closeErr := fd.Close()
+                if closeErr != nil && err == nil {
+                        err = errors.Annotate(closeErr).Reason("failed to Close").Err()
+                }
+        }()
+
+        // Emit a series of "Distribution==Version" strings.
+        seen := make(map[Name]struct{}, len(wheels))
+        for _, wheel := range wheels {
+                // Only mention a given Distribution/Version once.
+                archetype := Name{
+                        Distribution: wheel.Distribution,
+                        Version:      wheel.Version,
+                }
+                if _, ok := seen[archetype]; ok {
+                        // Already seen a package for this archetype, skip it.
+                        continue
+                }
+                seen[archetype] = struct{}{}
+
+                if _, err := fmt.Fprintf(fd, "%s==%s\n", archetype.Distribution, archetype.Version); err != nil {
+                        return errors.Annotate(err).Reason("failed to write to requirements file").Err()
+                }
+        }
+
+        return nil
+}