From f5530d9ed4e0afb4709a5ad28c5191e0ac18b64a Mon Sep 17 00:00:00 2001 From: "Denis Bilenko (aider)" Date: Mon, 3 Mar 2025 20:48:37 +0100 Subject: [PATCH] feat: Enhance wheel filename parsing and add target wheel existence check --- libs/patchwheel/patchwheel.go | 76 ++++++++++++++++++++++++++--- libs/patchwheel/patchwheel_test.go | 77 +++++++++++++++++++++++++----- 2 files changed, 134 insertions(+), 19 deletions(-) diff --git a/libs/patchwheel/patchwheel.go b/libs/patchwheel/patchwheel.go index 549f882de..d0c1b8297 100644 --- a/libs/patchwheel/patchwheel.go +++ b/libs/patchwheel/patchwheel.go @@ -125,30 +125,94 @@ func readFile(file *zip.File) ([]byte, error) { return io.ReadAll(rc) } +// WheelInfo contains information extracted from a wheel filename +type WheelInfo struct { + Distribution string // Package distribution name + Version string // Package version + Tags []string // Python tags (python_tag, abi_tag, platform_tag) +} + // ExtractVersionFromWheelFilename extracts the version from a wheel filename. // Wheel filenames follow the pattern: {distribution}-{version}-{python_tag}-{abi_tag}-{platform_tag}.whl func ExtractVersionFromWheelFilename(filename string) (string, error) { + info, err := ParseWheelFilename(filename) + if err != nil { + return "", err + } + return info.Version, nil +} + +// ParseWheelFilename parses a wheel filename and extracts its components. +// Wheel filenames follow the pattern: {distribution}-{version}-{python_tag}-{abi_tag}-{platform_tag}.whl +func ParseWheelFilename(filename string) (*WheelInfo, error) { base := filepath.Base(filename) parts := strings.Split(base, "-") if len(parts) < 5 || !strings.HasSuffix(parts[len(parts)-1], ".whl") { - return "", fmt.Errorf("invalid wheel filename format: %s", filename) + return nil, fmt.Errorf("invalid wheel filename format: %s", filename) } - - // If there are more than 5 parts, the distribution name might contain hyphens - // The version is always the second element from the end minus 3 (for the tags) - return parts[len(parts)-4], nil + + // The last three parts are always tags + tagStartIdx := len(parts) - 3 + + // Everything before the tags except the version is the distribution + versionIdx := tagStartIdx - 1 + + // Distribution may contain hyphens, so join all parts before the version + distribution := strings.Join(parts[:versionIdx], "-") + version := parts[versionIdx] + + // Extract tags (remove .whl from the last one) + tags := make([]string, 3) + copy(tags, parts[tagStartIdx:]) + tags[2] = strings.TrimSuffix(tags[2], ".whl") + + return &WheelInfo{ + Distribution: distribution, + Version: version, + Tags: tags, + }, nil } // PatchWheel patches a Python wheel file by updating its version in METADATA and RECORD. // It returns the path to the new wheel. // The function is idempotent: repeated calls with the same input will produce the same output. +// If the target wheel already exists, it returns the path to the existing wheel without processing. func PatchWheel(ctx context.Context, path, outputDir string) (string, error) { + // Get the modification time of the input wheel fileInfo, err := os.Stat(path) if err != nil { return "", err } wheelMtime := fileInfo.ModTime().UTC() - + + // Parse the wheel filename to extract components + wheelInfo, err := ParseWheelFilename(path) + if err != nil { + return "", err + } + + // Get the base version without any local version + baseVersion := strings.SplitN(wheelInfo.Version, "+", 2)[0] + + // Calculate the timestamp suffix for the new version + dt := strings.Replace(wheelMtime.Format("20060102150405.00"), ".", "", 1) + dt = strings.Replace(dt, ".", "", 1) + newVersion := baseVersion + "+" + dt + + // Create the new wheel filename + newFilename := fmt.Sprintf("%s-%s-%s.whl", + wheelInfo.Distribution, + newVersion, + strings.Join(wheelInfo.Tags, "-")) + outpath := filepath.Join(outputDir, newFilename) + + // Check if the target wheel already exists + if _, err := os.Stat(outpath); err == nil { + // Target wheel already exists, return its path + return outpath, nil + } + + // Target wheel doesn't exist, proceed with patching r, err := zip.OpenReader(path) if err != nil { return "", err diff --git a/libs/patchwheel/patchwheel_test.go b/libs/patchwheel/patchwheel_test.go index c1dfad041..7c3fd2952 100644 --- a/libs/patchwheel/patchwheel_test.go +++ b/libs/patchwheel/patchwheel_test.go @@ -106,28 +106,67 @@ func getWheel(t *testing.T, dir string) string { return matches[0] } -// TestExtractVersionFromWheelFilename tests the ExtractVersionFromWheelFilename function. -func TestExtractVersionFromWheelFilename(t *testing.T) { +// TestParseWheelFilename tests the ParseWheelFilename function. +func TestParseWheelFilename(t *testing.T) { tests := []struct { - filename string - wantVersion string - wantErr bool + filename string + wantDistribution string + wantVersion string + wantTags []string + wantErr bool }{ - {"myproj-0.1.0-py3-none-any.whl", "0.1.0", false}, - {"myproj-0.1.0+20240303123456-py3-none-any.whl", "0.1.0+20240303123456", false}, - {"my-proj-with-hyphens-0.1.0-py3-none-any.whl", "0.1.0", false}, - {"invalid-filename.txt", "", true}, - {"not-enough-parts-py3.whl", "", true}, + { + filename: "myproj-0.1.0-py3-none-any.whl", + wantDistribution: "myproj", + wantVersion: "0.1.0", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "myproj-0.1.0+20240303123456-py3-none-any.whl", + wantDistribution: "myproj", + wantVersion: "0.1.0+20240303123456", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "my-proj-with-hyphens-0.1.0-py3-none-any.whl", + wantDistribution: "my-proj-with-hyphens", + wantVersion: "0.1.0", + wantTags: []string{"py3", "none", "any"}, + wantErr: false, + }, + { + filename: "invalid-filename.txt", + wantDistribution: "", + wantVersion: "", + wantTags: nil, + wantErr: true, + }, + { + filename: "not-enough-parts-py3.whl", + wantDistribution: "", + wantVersion: "", + wantTags: nil, + wantErr: true, + }, } for _, tt := range tests { t.Run(tt.filename, func(t *testing.T) { - gotVersion, err := ExtractVersionFromWheelFilename(tt.filename) + info, err := ParseWheelFilename(tt.filename) if tt.wantErr { require.Error(t, err) } else { require.NoError(t, err) - require.Equal(t, tt.wantVersion, gotVersion) + require.Equal(t, tt.wantDistribution, info.Distribution) + require.Equal(t, tt.wantVersion, info.Version) + require.Equal(t, tt.wantTags, info.Tags) + + // Also test that ExtractVersionFromWheelFilename returns the same version + version, err := ExtractVersionFromWheelFilename(tt.filename) + require.NoError(t, err) + require.Equal(t, tt.wantVersion, version) } }) } @@ -155,14 +194,26 @@ func TestPatchWheel(t *testing.T) { origWheel := getWheel(t, distDir) // t.Logf("Found origWheel: %s", origWheel) + // First patch patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir) require.NoError(t, err) // t.Logf("origWheel=%s patchedWheel=%s", origWheel, patchedWheel) - + + // Get file info of the patched wheel + patchedInfo, err := os.Stat(patchedWheel) + require.NoError(t, err) + patchedTime := patchedInfo.ModTime() + // Test idempotency - patching the same wheel again should produce the same result + // and should not recreate the file (file modification time should remain the same) patchedWheel2, err := PatchWheel(context.Background(), origWheel, distDir) require.NoError(t, err) require.Equal(t, patchedWheel, patchedWheel2, "PatchWheel is not idempotent") + + // Check that the file wasn't recreated + patchedInfo2, err := os.Stat(patchedWheel2) + require.NoError(t, err) + require.Equal(t, patchedTime, patchedInfo2.ModTime(), "File was recreated when it shouldn't have been") runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel)