github.com/replit/upm@v0.0.0-20240423230255-9ce4fc3ea24c/internal/backends/python/gen_pypi_map/install_diff.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"os"
     7  	"os/exec"
     8  	"path/filepath"
     9  	"strings"
    10  	"time"
    11  )
    12  
    13  type InstallDiffResponse struct {
    14  	Modules []string `json:"modules"`
    15  	Error   string   `json:"error"`
    16  }
    17  
    18  /* discoverValidSuffixes probes the current runtime to discover what Python thinks
    19   * are importable.
    20   * This should make this broadly more portable, possibly even to Windows, or other
    21   * exciting runtimes.
    22   *
    23   * The constants here come from:
    24   *     https://github.com/python/cpython/blob/f59f90b5bccb9e7ac522bc779ab1f6bf11bb4aa3/Lib/modulefinder.py#L272-L274
    25   * and we should endeavor to keep them in sync.
    26   */
    27  func discoverValidSuffixes() ([]string, error) {
    28  	script := []string{
    29  		"from importlib.machinery import SOURCE_SUFFIXES, BYTECODE_SUFFIXES, EXTENSION_SUFFIXES",
    30  		"print('\\n'.join(SOURCE_SUFFIXES + BYTECODE_SUFFIXES + EXTENSION_SUFFIXES))",
    31  	}
    32  	cmd := exec.Command("python", "-c", strings.Join(script, "\n"))
    33  
    34  	var out bytes.Buffer
    35  	cmd.Stdout = &out
    36  
    37  	err := cmd.Run()
    38  	if err != nil {
    39  		return nil, PypiError{InstallFailure, "Installer failed", err}
    40  	}
    41  
    42  	res := strings.Split(strings.TrimSpace(out.String()), "\n")
    43  
    44  	return res, nil
    45  }
    46  
    47  func pipInstall(pkg string, root string, timeout time.Duration) error {
    48  	// Run pip to install just the package, so we can statically analyze it
    49  	cmd := exec.Command("pip", "install", "--no-deps", "--target", root, pkg)
    50  
    51  	if cmd == nil {
    52  		return PypiError{InstallFailure, "Failed to initialize exec.Command", nil}
    53  	}
    54  
    55  	killed := false
    56  	installing := true
    57  	go func() {
    58  		start := time.Now()
    59  		for installing {
    60  			elapsed := time.Since(start)
    61  			if elapsed > timeout {
    62  				err := cmd.Process.Signal(os.Interrupt)
    63  				if err == os.ErrProcessDone {
    64  					break
    65  				}
    66  				if err != nil {
    67  					err = cmd.Process.Signal(os.Kill)
    68  				}
    69  				if err == os.ErrProcessDone {
    70  					break
    71  				}
    72  				if err != nil {
    73  					fmt.Fprintf(os.Stderr, "Struggling to kill %d, %v\n", cmd.Process.Pid, err)
    74  				} else {
    75  					killed = true
    76  				}
    77  				break
    78  			}
    79  			time.Sleep(1 * time.Second)
    80  		}
    81  	}()
    82  
    83  	_, err := cmd.StdoutPipe()
    84  
    85  	if err != nil {
    86  		return PypiError{InstallFailure, "Failed to redirect stdout", err}
    87  	}
    88  
    89  	err = cmd.Run()
    90  	if err != nil {
    91  		if killed {
    92  			return PypiError{InstallFailure, "Exceeded timeout", err}
    93  		} else {
    94  			return PypiError{InstallFailure, "Failed to install", err}
    95  		}
    96  	}
    97  
    98  	installing = false
    99  	return nil
   100  }
   101  
   102  func InstallDiff(metadata PackageData, timeout time.Duration) ([]string, error) {
   103  	root := "/tmp/pypi/" + metadata.Info.Name
   104  
   105  	var err error
   106  	for attempts := 5; attempts > 0; attempts -= 1 {
   107  		_ = os.RemoveAll(root)
   108  		err = pipInstall(metadata.Info.Name, root, timeout)
   109  		if err == nil {
   110  			break
   111  		}
   112  	}
   113  
   114  	if err != nil {
   115  		return nil, err
   116  	}
   117  
   118  	suffixes, err := discoverValidSuffixes()
   119  	if err != nil {
   120  		return nil, err
   121  	}
   122  
   123  	var fewestSlashes *int
   124  	modules := make(map[string]bool)
   125  	err = filepath.WalkDir(root, func(fpath string, entry os.DirEntry, err error) error {
   126  		if err != nil {
   127  			return err
   128  		}
   129  		if entry.IsDir() {
   130  			// __pycache__ doesn't seem to have a constant, so we just presume it will
   131  			// never change.
   132  			if strings.HasSuffix(entry.Name(), ".dist-info") || entry.Name() == "__pycache__" {
   133  				return filepath.SkipDir
   134  			}
   135  		}
   136  
   137  		for _, suffix := range suffixes {
   138  			if !entry.IsDir() && strings.HasSuffix(entry.Name(), suffix) {
   139  				relpath, _ := filepath.Rel(root, fpath)
   140  				relpath = strings.TrimSuffix(relpath, suffix)
   141  				// Intention here is to trim "special" files, __init__.py and __main__.py, among others
   142  				if dir, last := filepath.Split(relpath); strings.HasPrefix(last, "__") && strings.HasSuffix(last, "__") {
   143  					relpath = dir
   144  				}
   145  				relpath = strings.TrimSuffix(relpath, string(filepath.Separator))
   146  				module := strings.ReplaceAll(relpath, string(filepath.Separator), ".")
   147  
   148  				// Do our best to find the lowest-common module root to avoid ballooning search space
   149  				// A good example of this is flask-mysql, which installs into flaskext.mysql
   150  				currentSlashes := strings.Count(relpath, string(filepath.Separator))
   151  				if fewestSlashes == nil || currentSlashes < *fewestSlashes {
   152  					fewestSlashes = &currentSlashes
   153  					modules = make(map[string]bool)
   154  					modules[module] = true
   155  				} else if currentSlashes == *fewestSlashes {
   156  					modules[module] = true
   157  				}
   158  			}
   159  		}
   160  
   161  		return nil
   162  	})
   163  
   164  	os.RemoveAll(root)
   165  
   166  	moduleList := []string{}
   167  	for module := range modules {
   168  		moduleList = append(moduleList, module)
   169  	}
   170  
   171  	return moduleList, err
   172  }