// 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 (

"context"

_ "embed"

"fmt"

"strings"

"github.com/emirpasic/gods/sets/treeset"

godsutils "github.com/emirpasic/gods/utils"

"golang.org/x/sync/errgroup"

)

// python3Parser implements a parser for Python files that extracts the modules

// as seen in the import statements.

type python3Parser struct {

// The value of language.GenerateArgs.Config.RepoRoot.

repoRoot string

// The value of language.GenerateArgs.Rel.

relPackagePath string

// The function that determines if a dependency is ignored from a Gazelle

// directive. It's the signature of pythonconfig.Config.IgnoresDependency.

ignoresDependency func(dep string) bool

}

// newPython3Parser constructs a new python3Parser.

func newPython3Parser(

repoRoot string,

relPackagePath string,

ignoresDependency func(dep string) bool,

) *python3Parser {

return &python3Parser{

repoRoot: repoRoot,

relPackagePath: relPackagePath,

ignoresDependency: ignoresDependency,

}

}

// parseSingle parses a single Python file and returns the extracted modules

// from the import statements as well as the parsed comments.

func (p *python3Parser) parseSingle(pyFilename string) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {

pyFilenames := treeset.NewWith(godsutils.StringComparator)

pyFilenames.Add(pyFilename)

return p.parse(pyFilenames)

}

// parse parses multiple Python files and returns the extracted modules from

// the import statements as well as the parsed comments.

func (p *python3Parser) parse(pyFilenames *treeset.Set) (*treeset.Set, map[string]*treeset.Set, *annotations, error) {

modules := treeset.NewWith(moduleComparator)

g, ctx := errgroup.WithContext(context.Background())

ch := make(chan struct{}, 6) // Limit the number of concurrent parses.

chRes := make(chan *ParserOutput, len(pyFilenames.Values()))

for _, v := range pyFilenames.Values() {

ch <- struct{}{}

g.Go(func(filename string) func() error {

return func() error {

defer func() {

<-ch

}()

res, err := NewFileParser().ParseFile(ctx, p.repoRoot, p.relPackagePath, filename)

if err != nil {

return err

}

chRes <- res

return nil

}

}(v.(string)))

}

if err := g.Wait(); err != nil {

return nil, nil, nil, err

}

close(ch)

close(chRes)

mainModules := make(map[string]*treeset.Set, len(chRes))

allAnnotations := new(annotations)

allAnnotations.ignore = make(map[string]struct{})

for res := range chRes {

if res.HasMain {

mainModules[res.FileName] = treeset.NewWith(moduleComparator)

}

annotations, err := annotationsFromComments(res.Comments)

if err != nil {

return nil, nil, nil, fmt.Errorf("failed to parse annotations: %w", err)

}

for _, m := range res.Modules {

// Check for ignored dependencies set via an annotation to the Python

// module.

if annotations.ignores(m.Name) || annotations.ignores(m.From) {

continue

}

// Check for ignored dependencies set via a Gazelle directive in a BUILD

// file.

if p.ignoresDependency(m.Name) || p.ignoresDependency(m.From) {

continue

}

modules.Add(m)

if res.HasMain {

mainModules[res.FileName].Add(m)

}

}

// Collect all annotations from each file into a single annotations struct.

for k, v := range annotations.ignore {

allAnnotations.ignore[k] = v

}

allAnnotations.includeDeps = append(allAnnotations.includeDeps, annotations.includeDeps...)

}

allAnnotations.includeDeps = removeDupesFromStringTreeSetSlice(allAnnotations.includeDeps)

return modules, mainModules, allAnnotations, nil

}

// removeDupesFromStringTreeSetSlice takes a []string, makes a set out of the

// elements, and then returns a new []string with all duplicates removed. Order

// is preserved.

func removeDupesFromStringTreeSetSlice(array []string) []string {

s := treeset.NewWith(godsutils.StringComparator)

for _, v := range array {

s.Add(v)

}

dedupe := make([]string, s.Size())

for i, v := range s.Values() {

dedupe[i] = fmt.Sprint(v)

}

return dedupe

}

// module represents a fully-qualified, dot-separated, Python module as seen on

// the import statement, alongside the line number where it happened.

type module struct {

// The fully-qualified, dot-separated, Python module name as seen on import

// statements.

Name string `json:"name"`

// The line number where the import happened.

LineNumber uint32 `json:"lineno"`

// The path to the module file relative to the Bazel workspace root.

Filepath string `json:"filepath"`

// If this was a from import, e.g. from foo import bar, From indicates the module

// from which it is imported.

From string `json:"from"`

}

// moduleComparator compares modules by name.

func moduleComparator(a, b interface{}) int {

return godsutils.StringComparator(a.(module).Name, b.(module).Name)

}

// annotationKind represents Gazelle annotation kinds.

type annotationKind string

const (

// The Gazelle annotation prefix.

annotationPrefix string = "gazelle:"

// The ignore annotation kind. E.g. '# gazelle:ignore <module_name>'.

annotationKindIgnore annotationKind = "ignore"

annotationKindIncludeDep annotationKind = "include_dep"

)

// comment represents a Python comment.

type comment string

// asAnnotation returns an annotation object if the comment has the

// annotationPrefix.

func (c *comment) asAnnotation() (*annotation, error) {

uncomment := strings.TrimLeft(string(*c), "# ")

if !strings.HasPrefix(uncomment, annotationPrefix) {

return nil, nil

}

withoutPrefix := strings.TrimPrefix(uncomment, annotationPrefix)

annotationParts := strings.SplitN(withoutPrefix, " ", 2)

if len(annotationParts) < 2 {

return nil, fmt.Errorf("`%s` requires a value", *c)

}

return &annotation{

kind: annotationKind(annotationParts[0]),

value: annotationParts[1],

}, nil

}

// annotation represents a single Gazelle annotation parsed from a Python

// comment.

type annotation struct {

kind annotationKind

value string

}

// annotations represent the collection of all Gazelle annotations parsed out of

// the comments of a Python module.

type annotations struct {

// The parsed modules to be ignored by Gazelle.

ignore map[string]struct{}

// Labels that Gazelle should include as deps of the generated target.

includeDeps []string

}

// annotationsFromComments returns all the annotations parsed out of the

// comments of a Python module.

func annotationsFromComments(comments []comment) (*annotations, error) {

ignore := make(map[string]struct{})

includeDeps := []string{}

for _, comment := range comments {

annotation, err := comment.asAnnotation()

if err != nil {

return nil, err

}

if annotation != nil {

if annotation.kind == annotationKindIgnore {

modules := strings.Split(annotation.value, ",")

for _, m := range modules {

if m == "" {

continue

}

m = strings.TrimSpace(m)

ignore[m] = struct{}{}

}

}

if annotation.kind == annotationKindIncludeDep {

targets := strings.Split(annotation.value, ",")

for _, t := range targets {

if t == "" {

continue

}

t = strings.TrimSpace(t)

includeDeps = append(includeDeps, t)

}

}

}

}

return &annotations{

ignore: ignore,

includeDeps: includeDeps,

}, nil

}

// ignored returns true if the given module was ignored via the ignore

// annotation.

func (a *annotations) ignores(module string) bool {

_, ignores := a.ignore[module]

return ignores

}