github.com/databricks/cli@v0.203.0/bundle/artifacts/whl/autodetect.go (about)

     1  package whl
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"path/filepath"
     8  	"regexp"
     9  	"time"
    10  
    11  	"github.com/databricks/cli/bundle"
    12  	"github.com/databricks/cli/bundle/config"
    13  	"github.com/databricks/cli/libs/cmdio"
    14  )
    15  
    16  type detectPkg struct {
    17  }
    18  
    19  func DetectPackage() bundle.Mutator {
    20  	return &detectPkg{}
    21  }
    22  
    23  func (m *detectPkg) Name() string {
    24  	return "artifacts.whl.AutoDetect"
    25  }
    26  
    27  func (m *detectPkg) Apply(ctx context.Context, b *bundle.Bundle) error {
    28  	cmdio.LogString(ctx, "artifacts.whl.AutoDetect: Detecting Python wheel project...")
    29  
    30  	// checking if there is setup.py in the bundle root
    31  	setupPy := filepath.Join(b.Config.Path, "setup.py")
    32  	_, err := os.Stat(setupPy)
    33  	if err != nil {
    34  		cmdio.LogString(ctx, "artifacts.whl.AutoDetect: No Python wheel project found at bundle root folder")
    35  		return nil
    36  	}
    37  
    38  	cmdio.LogString(ctx, fmt.Sprintf("artifacts.whl.AutoDetect: Found Python wheel project at %s", b.Config.Path))
    39  	module := extractModuleName(setupPy)
    40  
    41  	if b.Config.Artifacts == nil {
    42  		b.Config.Artifacts = make(map[string]*config.Artifact)
    43  	}
    44  
    45  	pkgPath, err := filepath.Abs(b.Config.Path)
    46  	if err != nil {
    47  		return err
    48  	}
    49  	b.Config.Artifacts[module] = &config.Artifact{
    50  		Path: pkgPath,
    51  		Type: config.ArtifactPythonWheel,
    52  	}
    53  
    54  	return nil
    55  }
    56  
    57  func extractModuleName(setupPy string) string {
    58  	bytes, err := os.ReadFile(setupPy)
    59  	if err != nil {
    60  		return randomName()
    61  	}
    62  
    63  	content := string(bytes)
    64  	r := regexp.MustCompile(`name=['"](.*)['"]`)
    65  	matches := r.FindStringSubmatch(content)
    66  	if len(matches) == 0 {
    67  		return randomName()
    68  	}
    69  	return matches[1]
    70  }
    71  
    72  func randomName() string {
    73  	return fmt.Sprintf("artifact%d", time.Now().Unix())
    74  }