Handle multiple wheel versions

pypi/BUILD

@@ -57,9 +57,9 @@ py_binary(
     data=[
         ":pip",
         ":wheeltool",
-
         "//tools:pip_generate",
-    ]
+    ],
+    visibility = ["//visibility:public"],
 )
 # checks tool path
 sh_test(

pypi/pip_generate.go

@@ -45,7 +45,7 @@ pyz_library(
 
 var pipLogLinkPattern = regexp.MustCompile(`^\s*(Found|Skipping) link\s*(http[^ #]+\.whl)`)
 
-var pyPIPlatforms = []struct {
+var platformDefs = []struct {
 	bazelPlatform string
 	// https://www.python.org/dev/peps/pep-0425/
 	pyPIPlatform string
@@ -211,7 +211,7 @@ type wheelToolOutput struct {
 	Extras   map[string][]string `json:"extras"`
 }
 
-func wheelDependencies(pythonPath string, wheelToolPath string, path string) ([]string, map[string][]string, error) {
+func wheelDependencies(pythonPath string, wheelToolPath string, path string, verbose bool) ([]string, map[string][]string, error) {
 	start := time.Now()
 	wheelToolProcess := exec.Command(wheelToolPath, path)
 	wheelToolProcess.Stderr = os.Stderr
@@ -220,7 +220,9 @@ func wheelDependencies(pythonPath string, wheelToolPath string, path string) ([]
 		return nil, nil, err
 	}
 	end := time.Now()
-	fmt.Printf("wheeltool %s took %s\n", filepath.Base(path), end.Sub(start).String())
+	if verbose {
+		fmt.Printf("wheeltool %s took %s\n", filepath.Base(path), end.Sub(start).String())
+	}
 	output := &wheelToolOutput{}
 	err = json.Unmarshal(outputBytes, output)
 	if err != nil {
@@ -230,9 +232,9 @@ func wheelDependencies(pythonPath string, wheelToolPath string, path string) ([]
 }
 
 func bazelPlatform(filename string) string {
-	for _, platformDefinition := range pyPIPlatforms {
-		if strings.Contains(filename, platformDefinition.pyPIPlatform) {
-			return platformDefinition.bazelPlatform
+	for _, platformDef := range platformDefs {
+		if strings.Contains(filename, platformDef.pyPIPlatform) {
+			return platformDef.bazelPlatform
 		}
 	}
 	return ""
@@ -300,6 +302,7 @@ func main() {
 		"Path to tool to output requirements from a wheel")
 	pythonPath := flag.String("pythonPath", "python", "Path to version of Python to use when running pip")
 	workspacePrefix := flag.String("workspacePrefix", "pypi_", "Prefix for generated repo rules")
+	shouldDeleteUnusedWheels := flag.Bool("deleteUnusedWheels", false, "Whether to delete wheels in `wheelDir` that are no longer used")
 	flag.Parse()
 	if *requirements == "" || *outputDir == "" {
 		fmt.Fprintln(os.Stderr, "Error: -requirements and -outputDir are required")
@@ -354,7 +357,7 @@ func main() {
 		panic(err)
 	}
 	pipProcess.Stderr = os.Stderr
-	fmt.Println("running pip to resolve dependencies ...")
+	fmt.Println("Running pip to resolve dependencies...")
 	if *verbose {
 		fmt.Printf("  command: %s %s\n", *pythonPath, strings.Join(pipProcess.Args, " "))
 	}
@@ -396,6 +399,7 @@ func main() {
 	pipEnd := time.Now()
 	fmt.Printf("pip executed in %v\n", pipEnd.Sub(pipStart).String())
 
+	fmt.Printf("Processing downloaded wheels...\n")
 	dirEntries, err := ioutil.ReadDir(tempDir)
 	if err != nil {
 		panic(err)
@@ -440,30 +444,42 @@ func main() {
 		bazelPlatform := bazelPlatform(entry.Name())
 		if bazelPlatform != "" {
 			// attempt to find all other platform wheels
-			matchedPlatforms := map[string]string{}
+			platformToWheelLink := map[string]string{}
 			matchPrefix := packageName + "-" + version + "-"
 			for wheelFile, link := range wheelFilenameToLink {
 				if strings.HasPrefix(wheelFile, matchPrefix) {
-					for _, pyPIPlatform := range pyPIPlatforms {
-						if pyPIPlatform.bazelPlatform == bazelPlatform {
+					for _, platformDef := range platformDefs {
+						if platformDef.bazelPlatform == bazelPlatform {
 							continue
 						}
-						if strings.Contains(wheelFile, pyPIPlatform.pyPIPlatform) {
-							if matchedPlatforms[pyPIPlatform.bazelPlatform] != "" {
-								panic("found duplicate wheels for platform")
+						if strings.Contains(wheelFile, platformDef.pyPIPlatform) {
+							existingWheelLink := platformToWheelLink[platformDef.bazelPlatform]
+							if existingWheelLink != "" {
+								// There are two versions. Need to pick one. For
+								// now, just pick alphabetically largest to ensure
+								// determinism.
+								fmt.Fprintf(os.Stderr, "Warning: two acceptable wheels found\n")
+								if link < existingWheelLink {
+									fmt.Fprintf(os.Stderr, "...picking %s instead of %s\n",
+										filepath.Base(existingWheelLink), filepath.Base(link))
+									link = existingWheelLink
+								} else {
+									fmt.Fprintf(os.Stderr, "...picking %s instead of %s\n",
+										filepath.Base(link), filepath.Base(existingWheelLink))
+								}
 							}
-							matchedPlatforms[pyPIPlatform.bazelPlatform] = link
+							platformToWheelLink[platformDef.bazelPlatform] = link
 						}
 					}
 				}
 			}
-			if len(matchedPlatforms)+1 != len(pyPIPlatforms) {
-				fmt.Fprintf(os.Stderr, "WARNING: could not find all platforms for %s; needs compilation?\n",
+			if len(platformToWheelLink)+1 != len(platformDefs) {
+				fmt.Fprintf(os.Stderr, "Warning: could not find all platformDefs for %s; needs compilation?\n",
 					entry.Name())
 			}
 
-			// download the other platforms and add info for those wheels
-			for _, link := range matchedPlatforms {
+			// download the other platformDefs and add info for those wheels
+			for _, link := range platformToWheelLink {
 				// download this PyPI wheel
 				filePart := filepath.Base(link)
 				destPath := path.Join(tempDir, filePart)
@@ -501,7 +517,7 @@ func main() {
 				panic(err)
 			}
 
-			deps, extras, err := wheelDependencies(*pythonPath, *wheelToolPath, partialInfo.filePath)
+			deps, extras, err := wheelDependencies(*pythonPath, *wheelToolPath, partialInfo.filePath, *verbose)
 			if err != nil {
 				panic(err)
 			}
@@ -563,7 +579,7 @@ func main() {
 
 		// ensure output is reproducible: output extras in the same order
 		extraNames := []string{}
-		for extraName, _ := range dependency.wheels[0].extras {
+		for extraName := range dependency.wheels[0].extras {
 			extraNames = append(extraNames, extraName)
 		}
 		sort.Strings(extraNames)
@@ -609,15 +625,48 @@ func main() {
 
 	// Lastly, make `http_file` repo rules for PyPI links.
 	fmt.Fprintln(outputBzlFile, "\ndef pypi_repositories():")
+	wroteAtLeastOne := false
 	// Don't call existing_rules repeatedly:
 	// https://github.com/bazelbuild/bazel/blob/master/site/docs/skylark/cookbook.md#aggregating-over-the-build-file
 	fmt.Fprintln(outputBzlFile, "    existing_rules = native.existing_rules()")
 	for _, dependency := range dependencies {
 		for _, wheel := range dependency.wheels {
 			if !wheel.useLocalWheel {
+				wroteAtLeastOne = true
 				name := wheel.bazelWorkspaceName(workspacePrefix)
 				fmt.Fprintf(outputBzlFile, wheel.makeBazelRule(&name, wheelDir))
 			}
 		}
 	}
+	if !wroteAtLeastOne {
+		fmt.Fprintln(outputBzlFile, "    pass")
+	}
+
+	if *shouldDeleteUnusedWheels {
+		deleteUnusedWheels(dependencies, path.Join(*outputDir, *wheelDir))
+	}
+
+	fmt.Printf("Done\n")
+}
+
+func deleteUnusedWheels(dependencies []pyPIDependency, absWheelDir string) {
+	// No, go does not have a `set` type. :/
+	wheelsThatShouldExist := map[string]bool{}
+	for _, dependency := range dependencies {
+		for _, wheel := range dependency.wheels {
+			wheelsThatShouldExist[wheel.filePath] = true
+		}
+	}
+
+	wheelPaths, err := filepath.Glob(path.Join(absWheelDir, "*"))
+	if err != nil {
+		panic(err)
+	}
+	for _, wheelPath := range wheelPaths {
+		_, present := wheelsThatShouldExist[wheelPath]
+		if !present {
+			fmt.Fprintf(os.Stderr, "Deleting unused wheel: %s\n", wheelPath)
+			os.Remove(wheelPath)
+		}
+	}
 }