go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/vpython/wheels/requirements.go (about)

     1  // Copyright 2024 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package wheels
    16  
    17  import (
    18  	"fmt"
    19  	"os"
    20  	"path/filepath"
    21  	"strings"
    22  
    23  	"go.chromium.org/luci/common/errors"
    24  )
    25  
    26  // WheelName is a parsed Python wheel name, defined here:
    27  // https://www.python.org/dev/peps/pep-0427/#file-name-convention
    28  //
    29  // {distribution}-{version}(-{build tag})?-{python tag}-{abi tag}-\
    30  // {platform tag}.whl .
    31  type wheelName struct {
    32  	Distribution string
    33  	Version      string
    34  	BuildTag     string
    35  	PythonTag    string
    36  	ABITag       string
    37  	PlatformTag  string
    38  }
    39  
    40  func (wn *wheelName) String() string {
    41  	parts := make([]string, 0, 6)
    42  	parts = append(parts, []string{
    43  		wn.Distribution,
    44  		wn.Version,
    45  	}...)
    46  	if wn.BuildTag != "" {
    47  		parts = append(parts, wn.BuildTag)
    48  	}
    49  	parts = append(parts, []string{
    50  		wn.PythonTag,
    51  		wn.ABITag,
    52  		wn.PlatformTag,
    53  	}...)
    54  	return strings.Join(parts, "-") + ".whl"
    55  }
    56  
    57  // ParseName parses a wheel Name from its filename.
    58  func parseName(v string) (wn wheelName, err error) {
    59  	base := strings.TrimSuffix(v, ".whl")
    60  	if len(base) == len(v) {
    61  		err = errors.New("missing .whl suffix")
    62  		return
    63  	}
    64  
    65  	skip := 0
    66  	switch parts := strings.Split(base, "-"); len(parts) {
    67  	case 6:
    68  		// Extra part: build tag.
    69  		wn.BuildTag = parts[2]
    70  		skip = 1
    71  		fallthrough
    72  
    73  	case 5:
    74  		wn.Distribution = parts[0]
    75  		wn.Version = parts[1]
    76  		wn.PythonTag = parts[2+skip]
    77  		wn.ABITag = parts[3+skip]
    78  		wn.PlatformTag = parts[4+skip]
    79  
    80  	default:
    81  		err = errors.Reason("unknown number of segments (%d)", len(parts)).Err()
    82  		return
    83  	}
    84  	return
    85  }
    86  
    87  // ScanDir identifies all wheel files in the immediate directory dir and
    88  // returns their parsed wheel names.
    89  func scanDir(dir string) ([]wheelName, error) {
    90  	globPattern := filepath.Join(dir, "*.whl")
    91  	matches, err := filepath.Glob(globPattern)
    92  	if err != nil {
    93  		return nil, errors.Annotate(err, "failed to list wheel directory: %s", globPattern).Err()
    94  	}
    95  
    96  	names := make([]wheelName, 0, len(matches))
    97  	for _, match := range matches {
    98  		switch st, err := os.Stat(match); {
    99  		case err != nil:
   100  			return nil, errors.Annotate(err, "failed to stat wheel in dir %s: %s", dir, match).Err()
   101  
   102  		case st.IsDir():
   103  			// Ignore directories.
   104  			continue
   105  
   106  		default:
   107  			// A ".whl" file.
   108  			name := filepath.Base(match)
   109  			wheelName, err := parseName(name)
   110  			if err != nil {
   111  				return nil, errors.Annotate(err, "failed to parse wheel from %s: %s", dir, name).Err()
   112  			}
   113  			names = append(names, wheelName)
   114  		}
   115  	}
   116  	return names, nil
   117  }
   118  
   119  // WriteRequirementsFile writes a valid "requirements.txt"-style pip
   120  // requirements file containing the supplied wheels.
   121  //
   122  // The generated requirements will request the exact wheel senver version (using
   123  // "==").
   124  func writeRequirementsFile(path string, wheels []wheelName) (err error) {
   125  	fd, err := os.Create(path)
   126  	if err != nil {
   127  		return errors.Annotate(err, "failed to create requirements file").Err()
   128  	}
   129  	defer func() {
   130  		closeErr := fd.Close()
   131  		if closeErr != nil && err == nil {
   132  			err = errors.Annotate(closeErr, "failed to Close").Err()
   133  		}
   134  	}()
   135  
   136  	// Emit a series of "Distribution==Version" strings.
   137  	seen := make(map[wheelName]struct{}, len(wheels))
   138  	for _, wheel := range wheels {
   139  		// Only mention a given Distribution/Version once.
   140  		archetype := wheelName{
   141  			Distribution: wheel.Distribution,
   142  			Version:      wheel.Version,
   143  		}
   144  		if _, ok := seen[archetype]; ok {
   145  			// Already seen a package for this archetype, skip it.
   146  			continue
   147  		}
   148  		seen[archetype] = struct{}{}
   149  
   150  		if _, err := fmt.Fprintf(fd, "%s==%s\n", archetype.Distribution, archetype.Version); err != nil {
   151  			return errors.Annotate(err, "failed to write to requirements file").Err()
   152  		}
   153  	}
   154  
   155  	return nil
   156  }