github.phpd.cn/thought-machine/please@v12.2.0+incompatible/tools/please_pex/pex/pex.go (about)

     1  // Package pex implements construction of .pex files in Go.
     2  // For performance reasons we've ultimately abandoned doing this in Python;
     3  // we were ultimately not using pex for much at construction time and
     4  // we already have most of what we need in Go via jarcat.
     5  package pex
     6  
     7  import (
     8  	"bytes"
     9  	"os"
    10  	"path"
    11  	"strings"
    12  
    13  	"tools/jarcat/zip"
    14  )
    15  
    16  // A Writer implements writing a .pex file in various steps.
    17  type Writer struct {
    18  	zipSafe                    bool
    19  	shebang                    string
    20  	realEntryPoint             string
    21  	testSrcs                   []string
    22  	testIncludes, testExcludes []string
    23  	testRunner                 string
    24  }
    25  
    26  // NewWriter constructs a new Writer.
    27  func NewWriter(entryPoint, interpreter string, zipSafe bool) *Writer {
    28  	pw := &Writer{
    29  		zipSafe:        zipSafe,
    30  		realEntryPoint: toPythonPath(entryPoint),
    31  	}
    32  	pw.SetShebang(interpreter)
    33  	return pw
    34  }
    35  
    36  // SetShebang sets the leading shebang that will be written to the file.
    37  func (pw *Writer) SetShebang(shebang string) {
    38  	if !path.IsAbs(shebang) {
    39  		shebang = "/usr/bin/env " + shebang
    40  	}
    41  	if !strings.HasPrefix(shebang, "#") {
    42  		shebang = "#!" + shebang
    43  	}
    44  	pw.shebang = shebang + "\n"
    45  }
    46  
    47  // SetTest sets this Writer to write tests using the given sources.
    48  // This overrides the entry point given earlier.
    49  func (pw *Writer) SetTest(srcs []string, usePyTest bool) {
    50  	pw.realEntryPoint = "test_main"
    51  	pw.testSrcs = srcs
    52  	if usePyTest {
    53  		// We only need xmlrunner for unittest, the equivalent is builtin to pytest.
    54  		pw.testExcludes = []string{".bootstrap/xmlrunner/*"}
    55  		pw.testRunner = "pytest.py"
    56  	} else {
    57  		pw.testIncludes = []string{
    58  			".bootstrap/xmlrunner",
    59  			".bootstrap/coverage",
    60  			".bootstrap/__init__.py",
    61  			".bootstrap/six.py",
    62  		}
    63  		pw.testRunner = "unittest.py"
    64  	}
    65  }
    66  
    67  // Write writes the pex to the given output file.
    68  func (pw *Writer) Write(out, moduleDir string) error {
    69  	f := zip.NewFile(out, true)
    70  	defer f.Close()
    71  
    72  	// Write preamble (i.e. the shebang that makes it executable)
    73  	if err := f.WritePreamble([]byte(pw.shebang)); err != nil {
    74  		return err
    75  	}
    76  
    77  	// Write required pex stuff for tests. Note that this executable is also a zipfile and we can
    78  	// jarcat it directly in (nifty, huh?).
    79  	if len(pw.testSrcs) != 0 {
    80  		f.Include = pw.testIncludes
    81  		f.Exclude = pw.testExcludes
    82  		if err := f.AddZipFile(os.Args[0]); err != nil {
    83  			return err
    84  		}
    85  	}
    86  
    87  	// Always write pex_main.py, with some templating.
    88  	b := MustAsset("pex_main.py")
    89  	b = bytes.Replace(b, []byte("__MODULE_DIR__"), []byte(moduleDir), 1)
    90  	b = bytes.Replace(b, []byte("__ENTRY_POINT__"), []byte(pw.realEntryPoint), 1)
    91  	b = bytes.Replace(b, []byte("__ZIP_SAFE__"), []byte(pythonBool(pw.zipSafe)), 1)
    92  
    93  	if len(pw.testSrcs) != 0 {
    94  		// If we're writing a test, we append test_main.py to it.
    95  		b2 := MustAsset("test_main.py")
    96  		b2 = bytes.Replace(b2, []byte("__TEST_NAMES__"), []byte(strings.Join(pw.testSrcs, ",")), 1)
    97  		b = append(b, b2...)
    98  		// It also needs an appropriate test runner.
    99  		b = append(b, MustAsset(pw.testRunner)...)
   100  	}
   101  	// We always append the final if __name__ == '__main__' bit.
   102  	b = append(b, MustAsset("pex_run.py")...)
   103  	return f.WriteFile("__main__.py", b, 0644)
   104  }
   105  
   106  // pythonBool returns a Python bool representation of a Go bool.
   107  func pythonBool(b bool) string {
   108  	if b {
   109  		return "True"
   110  	}
   111  	return "False"
   112  }
   113  
   114  // toPythonPath converts a normal path to a Python import path.
   115  func toPythonPath(p string) string {
   116  	ext := path.Ext(p)
   117  	return strings.Replace(p[:len(p)-len(ext)], "/", ".", -1)
   118  }