mirror of https://github.com/databricks/cli.git
feat: Enhance wheel filename parsing and add target wheel existence check
This commit is contained in:
parent
43be8ea727
commit
f5530d9ed4
|
@ -125,30 +125,94 @@ func readFile(file *zip.File) ([]byte, error) {
|
||||||
return io.ReadAll(rc)
|
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.
|
// ExtractVersionFromWheelFilename extracts the version from a wheel filename.
|
||||||
// Wheel filenames follow the pattern: {distribution}-{version}-{python_tag}-{abi_tag}-{platform_tag}.whl
|
// Wheel filenames follow the pattern: {distribution}-{version}-{python_tag}-{abi_tag}-{platform_tag}.whl
|
||||||
func ExtractVersionFromWheelFilename(filename string) (string, error) {
|
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)
|
base := filepath.Base(filename)
|
||||||
parts := strings.Split(base, "-")
|
parts := strings.Split(base, "-")
|
||||||
if len(parts) < 5 || !strings.HasSuffix(parts[len(parts)-1], ".whl") {
|
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 last three parts are always tags
|
||||||
// The version is always the second element from the end minus 3 (for the tags)
|
tagStartIdx := len(parts) - 3
|
||||||
return parts[len(parts)-4], nil
|
|
||||||
|
// 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.
|
// PatchWheel patches a Python wheel file by updating its version in METADATA and RECORD.
|
||||||
// It returns the path to the new wheel.
|
// It returns the path to the new wheel.
|
||||||
// The function is idempotent: repeated calls with the same input will produce the same output.
|
// 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) {
|
func PatchWheel(ctx context.Context, path, outputDir string) (string, error) {
|
||||||
|
// Get the modification time of the input wheel
|
||||||
fileInfo, err := os.Stat(path)
|
fileInfo, err := os.Stat(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
wheelMtime := fileInfo.ModTime().UTC()
|
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)
|
r, err := zip.OpenReader(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
|
|
|
@ -106,28 +106,67 @@ func getWheel(t *testing.T, dir string) string {
|
||||||
return matches[0]
|
return matches[0]
|
||||||
}
|
}
|
||||||
|
|
||||||
// TestExtractVersionFromWheelFilename tests the ExtractVersionFromWheelFilename function.
|
// TestParseWheelFilename tests the ParseWheelFilename function.
|
||||||
func TestExtractVersionFromWheelFilename(t *testing.T) {
|
func TestParseWheelFilename(t *testing.T) {
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
filename string
|
filename string
|
||||||
|
wantDistribution string
|
||||||
wantVersion string
|
wantVersion string
|
||||||
|
wantTags []string
|
||||||
wantErr bool
|
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},
|
filename: "myproj-0.1.0-py3-none-any.whl",
|
||||||
{"my-proj-with-hyphens-0.1.0-py3-none-any.whl", "0.1.0", false},
|
wantDistribution: "myproj",
|
||||||
{"invalid-filename.txt", "", true},
|
wantVersion: "0.1.0",
|
||||||
{"not-enough-parts-py3.whl", "", true},
|
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 {
|
for _, tt := range tests {
|
||||||
t.Run(tt.filename, func(t *testing.T) {
|
t.Run(tt.filename, func(t *testing.T) {
|
||||||
gotVersion, err := ExtractVersionFromWheelFilename(tt.filename)
|
info, err := ParseWheelFilename(tt.filename)
|
||||||
if tt.wantErr {
|
if tt.wantErr {
|
||||||
require.Error(t, err)
|
require.Error(t, err)
|
||||||
} else {
|
} else {
|
||||||
require.NoError(t, err)
|
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,15 +194,27 @@ func TestPatchWheel(t *testing.T) {
|
||||||
origWheel := getWheel(t, distDir)
|
origWheel := getWheel(t, distDir)
|
||||||
// t.Logf("Found origWheel: %s", origWheel)
|
// t.Logf("Found origWheel: %s", origWheel)
|
||||||
|
|
||||||
|
// First patch
|
||||||
patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir)
|
patchedWheel, err := PatchWheel(context.Background(), origWheel, distDir)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
// t.Logf("origWheel=%s patchedWheel=%s", origWheel, patchedWheel)
|
// 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
|
// 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)
|
patchedWheel2, err := PatchWheel(context.Background(), origWheel, distDir)
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
require.Equal(t, patchedWheel, patchedWheel2, "PatchWheel is not idempotent")
|
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)
|
runCmd(t, tempDir, "uv", "pip", "install", "-q", patchedWheel)
|
||||||
|
|
||||||
pyExec := filepath.Join(tempDir, ".venv", "bin", "python") // XXX Windows
|
pyExec := filepath.Join(tempDir, ".venv", "bin", "python") // XXX Windows
|
||||||
|
|
Loading…
Reference in New Issue