gazelle/python/resolve.go

// Copyright 2023 The Bazel Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package python

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/bazelbuild/bazel-gazelle/config"
	"github.com/bazelbuild/bazel-gazelle/label"
	"github.com/bazelbuild/bazel-gazelle/repo"
	"github.com/bazelbuild/bazel-gazelle/resolve"
	"github.com/bazelbuild/bazel-gazelle/rule"
	bzl "github.com/bazelbuild/buildtools/build"
	"github.com/emirpasic/gods/sets/treeset"
	godsutils "github.com/emirpasic/gods/utils"

	"github.com/bazel-contrib/rules_python/gazelle/pythonconfig"
)

const languageName = "py"

const (
	// resolvedDepsKey is the attribute key used to pass dependencies that don't
	// need to be resolved by the dependency resolver in the Resolver step.
	resolvedDepsKey = "_gazelle_python_resolved_deps"
)

// Resolver satisfies the resolve.Resolver interface. It resolves dependencies
// in rules generated by this extension.
type Resolver struct{}

// Name returns the name of the language. This is the prefix of the kinds of
// rules generated. E.g. py_library and py_binary.
func (*Resolver) Name() string { return languageName }

// Imports returns a list of ImportSpecs that can be used to import the rule
// r. This is used to populate RuleIndex.
//
// If nil is returned, the rule will not be indexed. If any non-nil slice is
// returned, including an empty slice, the rule will be indexed.
func (py *Resolver) Imports(c *config.Config, r *rule.Rule, f *rule.File) []resolve.ImportSpec {
	cfgs := c.Exts[languageName].(pythonconfig.Configs)
	cfg := cfgs[f.Pkg]
	srcs := r.AttrStrings("srcs")
	provides := make([]resolve.ImportSpec, 0, len(srcs)+1)
	for _, src := range srcs {
		ext := filepath.Ext(src)
		if ext != ".py" {
			continue
		}
		if cfg.PerFileGeneration() && len(srcs) > 1 && src == pyLibraryEntrypointFilename {
			// Do not provide import spec from __init__.py when it is being included as
			// part of another module.
			continue
		}
		pythonProjectRoot := cfg.PythonProjectRoot()
		provide := importSpecFromSrc(pythonProjectRoot, f.Pkg, src)
		provides = append(provides, provide)
	}
	if len(provides) == 0 {
		return nil
	}
	return provides
}

// importSpecFromSrc determines the ImportSpec based on the target that contains the src so that
// the target can be indexed for import statements that match the calculated src relative to the its
// Python project root.
func importSpecFromSrc(pythonProjectRoot, bzlPkg, src string) resolve.ImportSpec {
	pythonPkgDir := filepath.Join(bzlPkg, filepath.Dir(src))
	relPythonPkgDir, err := filepath.Rel(pythonProjectRoot, pythonPkgDir)
	if err != nil {
		panic(fmt.Errorf("unexpected failure: %v", err))
	}
	if relPythonPkgDir == "." {
		relPythonPkgDir = ""
	}
	pythonPkg := strings.ReplaceAll(relPythonPkgDir, "/", ".")
	filename := filepath.Base(src)
	if filename == pyLibraryEntrypointFilename {
		if pythonPkg != "" {
			return resolve.ImportSpec{
				Lang: languageName,
				Imp:  pythonPkg,
			}
		}
	}
	moduleName := strings.TrimSuffix(filename, ".py")
	var imp string
	if pythonPkg == "" {
		imp = moduleName
	} else {
		imp = fmt.Sprintf("%s.%s", pythonPkg, moduleName)
	}
	return resolve.ImportSpec{
		Lang: languageName,
		Imp:  imp,
	}
}

// Embeds returns a list of labels of rules that the given rule embeds. If
// a rule is embedded by another importable rule of the same language, only
// the embedding rule will be indexed. The embedding rule will inherit
// the imports of the embedded rule.
func (py *Resolver) Embeds(r *rule.Rule, from label.Label) []label.Label {
	// TODO(f0rmiga): implement.
	return make([]label.Label, 0)
}

// addDependency adds a dependency to either the regular deps or pyiDeps set based on
// whether the module is type-checking only. If a module is added as both
// non-type-checking and type-checking, it should end up in deps and not pyiDeps.
func addDependency(dep string, typeCheckingOnly bool, deps, pyiDeps *treeset.Set) {
	if typeCheckingOnly {
		if !deps.Contains(dep) {
			pyiDeps.Add(dep)
		}
	} else {
		deps.Add(dep)
		pyiDeps.Remove(dep)
	}
}

// Resolve translates imported libraries for a given rule into Bazel
// dependencies. Information about imported libraries is returned for each
// rule generated by language.GenerateRules in
// language.GenerateResult.Imports. Resolve generates a "deps" attribute (or
// the appropriate language-specific equivalent) for each import according to
// language-specific rules and heuristics.
func (py *Resolver) Resolve(
	c *config.Config,
	ix *resolve.RuleIndex,
	rc *repo.RemoteCache,
	r *rule.Rule,
	modulesRaw interface{},
	from label.Label,
) {
	// TODO(f0rmiga): may need to be defensive here once this Gazelle extension
	// join with the main Gazelle binary with other rules. It may conflict with
	// other generators that generate py_* targets.
	deps := treeset.NewWith(godsutils.StringComparator)
	pyiDeps := treeset.NewWith(godsutils.StringComparator)
	cfgs := c.Exts[languageName].(pythonconfig.Configs)
	cfg := cfgs[from.Pkg]

	if modulesRaw != nil {
		pythonProjectRoot := cfg.PythonProjectRoot()
		modules := modulesRaw.(*treeset.Set)
		it := modules.Iterator()
		explainDependency := os.Getenv("EXPLAIN_DEPENDENCY")
		hasFatalError := false
	MODULES_LOOP:
		for it.Next() {
			mod := it.Value().(Module)
			moduleName := mod.Name
			// Transform relative imports `.` or `..foo.bar` into the package path from root.
			if strings.HasPrefix(mod.From, ".") {
				if !cfg.ExperimentalAllowRelativeImports() {
					continue MODULES_LOOP
				}

				// Count number of leading dots in mod.From (e.g., ".." = 2, "...foo.bar" = 3)
				relativeDepth := strings.IndexFunc(mod.From, func(r rune) bool { return r != '.' })
				if relativeDepth == -1 {
					relativeDepth = len(mod.From)
				}

				// Extract final symbol (e.g., "some_function") from mod.Name
				imported := mod.Name
				if idx := strings.LastIndex(mod.Name, "."); idx >= 0 {
					imported = mod.Name[idx+1:]
				}

				// Optional subpath in 'from' clause, e.g. "from ...my_library.foo import x"
				fromPath := strings.TrimLeft(mod.From, ".")
				var fromParts []string
				if fromPath != "" {
					fromParts = strings.Split(fromPath, ".")
				}

				// Current Bazel package as path segments
				pkgParts := strings.Split(from.Pkg, "/")

				if relativeDepth-1 > len(pkgParts) {
					log.Printf("ERROR: Invalid relative import %q in %q: exceeds package root.", mod.Name, mod.Filepath)
					continue MODULES_LOOP
				}

				// Go up relativeDepth - 1 levels
				baseParts := pkgParts
				if relativeDepth > 1 {
					baseParts = pkgParts[:len(pkgParts)-(relativeDepth-1)]
				}
				// Build absolute module path
				absParts := append([]string{}, baseParts...) // base path
				absParts = append(absParts, fromParts...)    // subpath from 'from'
				absParts = append(absParts, imported)        // actual imported symbol

				moduleName = strings.Join(absParts, ".")
			}

			moduleParts := strings.Split(moduleName, ".")
			possibleModules := []string{moduleName}
			for len(moduleParts) > 1 {
				// Iterate back through the possible imports until
				// a match is found.
				// For example, "from foo.bar import baz" where baz is a module, we should try `foo.bar.baz` first, then
				// `foo.bar`, then `foo`.
				// In the first case, the import could be file `baz.py` in the directory `foo/bar`.
				// Or, the import could be variable `baz` in file `foo/bar.py`.
				// The import could also be from a standard module, e.g. `six.moves`, where
				// the dependency is actually `six`.
				moduleParts = moduleParts[:len(moduleParts)-1]
				possibleModules = append(possibleModules, strings.Join(moduleParts, "."))
			}
			errs := []error{}
		POSSIBLE_MODULE_LOOP:
			for _, moduleName := range possibleModules {
				imp := resolve.ImportSpec{Lang: languageName, Imp: moduleName}
				if override, ok := resolve.FindRuleWithOverride(c, imp, languageName); ok {
					if override.Repo == "" {
						override.Repo = from.Repo
					}
					if !override.Equal(from) {
						if override.Repo == from.Repo {
							override.Repo = ""
						}
						dep := override.Rel(from.Repo, from.Pkg).String()
						addDependency(dep, mod.TypeCheckingOnly, deps, pyiDeps)
						if explainDependency == dep {
							log.Printf("Explaining dependency (%s): "+
								"in the target %q, the file %q imports %q at line %d, "+
								"which resolves using the \"gazelle:resolve\" directive.\n",
								explainDependency, from.String(), mod.Filepath, moduleName, mod.LineNumber)
						}
						continue MODULES_LOOP
					}
				} else {
					if dep, distributionName, ok := cfg.FindThirdPartyDependency(moduleName); ok {
						addDependency(dep, mod.TypeCheckingOnly, deps, pyiDeps)
						// Add the type and stub dependencies if they exist.
						modules := []string{
							fmt.Sprintf("%s_stubs", strings.ToLower(distributionName)),
							fmt.Sprintf("%s_types", strings.ToLower(distributionName)),
							fmt.Sprintf("types_%s", strings.ToLower(distributionName)),
							fmt.Sprintf("stubs_%s", strings.ToLower(distributionName)),
						}
						for _, module := range modules {
							if dep, _, ok := cfg.FindThirdPartyDependency(module); ok {
								// Type stub packages are added as type-checking only.
								addDependency(dep, true, deps, pyiDeps)
							}
						}
						if explainDependency == dep {
							log.Printf("Explaining dependency (%s): "+
								"in the target %q, the file %q imports %q at line %d, "+
								"which resolves from the third-party module %q from the wheel %q.\n",
								explainDependency, from.String(), mod.Filepath, moduleName, mod.LineNumber, mod.Name, dep)
						}
						continue MODULES_LOOP
					} else {
						matches := ix.FindRulesByImportWithConfig(c, imp, languageName)
						if len(matches) == 0 {
							// Check if the imported module is part of the standard library.
							if isStdModule(Module{Name: moduleName}) {
								continue MODULES_LOOP
							} else if cfg.ValidateImportStatements() {
								err := fmt.Errorf(
									"%[1]q, line %[2]d: %[3]q is an invalid dependency: possible solutions:\n"+
										"\t1. Add it as a dependency in the requirements.txt file.\n"+
										"\t2. Use the '# gazelle:resolve py %[3]s TARGET_LABEL' BUILD file directive to resolve to a known dependency.\n"+
										"\t3. Ignore it with a comment '# gazelle:ignore %[3]s' in the Python file.\n",
									mod.Filepath, mod.LineNumber, moduleName,
								)
								errs = append(errs, err)
								continue POSSIBLE_MODULE_LOOP
							}
						}
						filteredMatches := make([]resolve.FindResult, 0, len(matches))
						for _, match := range matches {
							if match.IsSelfImport(from) {
								// Prevent from adding itself as a dependency.
								continue MODULES_LOOP
							}
							filteredMatches = append(filteredMatches, match)
						}
						if len(filteredMatches) == 0 {
							continue POSSIBLE_MODULE_LOOP
						}
						if len(filteredMatches) > 1 {
							sameRootMatches := make([]resolve.FindResult, 0, len(filteredMatches))
							for _, match := range filteredMatches {
								if strings.HasPrefix(match.Label.Pkg, pythonProjectRoot) {
									sameRootMatches = append(sameRootMatches, match)
								}
							}
							if len(sameRootMatches) != 1 {
								err := fmt.Errorf(
									"%[1]q, line %[2]d: multiple targets (%[3]s) may be imported with %[4]q: possible solutions:\n"+
										"\t1. Disambiguate the above multiple targets by removing duplicate srcs entries.\n"+
										"\t2. Use the '# gazelle:resolve py %[4]s TARGET_LABEL' BUILD file directive to resolve to one of the above targets.\n",
									mod.Filepath, mod.LineNumber, targetListFromResults(filteredMatches), moduleName)
								errs = append(errs, err)
								continue POSSIBLE_MODULE_LOOP
							}
							filteredMatches = sameRootMatches
						}
						matchLabel := filteredMatches[0].Label.Rel(from.Repo, from.Pkg)
						dep := matchLabel.String()
						addDependency(dep, mod.TypeCheckingOnly, deps, pyiDeps)
						if explainDependency == dep {
							log.Printf("Explaining dependency (%s): "+
								"in the target %q, the file %q imports %q at line %d, "+
								"which resolves from the first-party indexed labels.\n",
								explainDependency, from.String(), mod.Filepath, moduleName, mod.LineNumber)
						}
						continue MODULES_LOOP
					}
				}
			} // End possible modules loop.
			if len(errs) > 0 {
				// If, after trying all possible modules, we still haven't found anything, error out.
				joinedErrs := ""
				for _, err := range errs {
					joinedErrs = fmt.Sprintf("%s%s\n", joinedErrs, err)
				}
				log.Printf("ERROR: failed to validate dependencies for target %q:\n\n%v", from.String(), joinedErrs)
				hasFatalError = true
			}
		}
		if hasFatalError {
			os.Exit(1)
		}
	}

	addResolvedDeps(r, deps)

	if cfg.GeneratePyiDeps() {
		if !deps.Empty() {
			r.SetAttr("deps", convertDependencySetToExpr(deps))
		}
		if !pyiDeps.Empty() {
			r.SetAttr("pyi_deps", convertDependencySetToExpr(pyiDeps))
		}
	} else {
		// When generate_pyi_deps is false, merge both deps and pyiDeps into deps
		combinedDeps := treeset.NewWith(godsutils.StringComparator)
		combinedDeps.Add(deps.Values()...)
		combinedDeps.Add(pyiDeps.Values()...)

		if !combinedDeps.Empty() {
			r.SetAttr("deps", convertDependencySetToExpr(combinedDeps))
		}
	}
}

// addResolvedDeps adds the pre-resolved dependencies from the rule's private attributes
// to the provided deps set.
func addResolvedDeps(
	r *rule.Rule,
	deps *treeset.Set,
) {
	resolvedDeps := r.PrivateAttr(resolvedDepsKey).(*treeset.Set)
	if !resolvedDeps.Empty() {
		it := resolvedDeps.Iterator()
		for it.Next() {
			deps.Add(it.Value())
		}
	}
}

// targetListFromResults returns a string with the human-readable list of
// targets contained in the given results.
func targetListFromResults(results []resolve.FindResult) string {
	list := make([]string, len(results))
	for i, result := range results {
		list[i] = result.Label.String()
	}
	return strings.Join(list, ", ")
}

// convertDependencySetToExpr converts the given set of dependencies to an
// expression to be used in the deps attribute.
func convertDependencySetToExpr(set *treeset.Set) bzl.Expr {
	deps := make([]bzl.Expr, set.Size())
	it := set.Iterator()
	for it.Next() {
		dep := it.Value().(string)
		deps[it.Index()] = &bzl.StringExpr{Value: dep}
	}
	return &bzl.ListExpr{List: deps}
}