github.com/tiagovtristao/plz@v13.4.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  	"github.com/thought-machine/please/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   []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, testRunner string) {
    50  	pw.realEntryPoint = "test_main"
    51  	pw.testSrcs = srcs
    52  
    53  	commonIncludes := []string{
    54  		".bootstrap/coverage",
    55  		".bootstrap/__init__.py",
    56  		".bootstrap/six.py",
    57  	}
    58  
    59  	if testRunner == "pytest" {
    60  		// We only need xmlrunner for unittest, the equivalent is builtin to pytest.
    61  		pw.testIncludes = append(commonIncludes,
    62  			".bootstrap/pytest.py",
    63  			".bootstrap/_pytest",
    64  			".bootstrap/py",
    65  			".bootstrap/pluggy",
    66  			".bootstrap/attr",
    67  			".bootstrap/funcsigs",
    68  			".bootstrap/pkg_resources",
    69  		)
    70  
    71  		pw.testRunner = "pytest.py"
    72  	} else if testRunner == "behave" {
    73  		pw.testIncludes = append(commonIncludes,
    74  			".bootstrap/behave",
    75  			".bootstrap/parse.py",
    76  			".bootstrap/parse_type",
    77  			".bootstrap/traceback2",
    78  			".bootstrap/enum",
    79  			".bootstrap/win_unicode_console",
    80  			".bootstrap/colorama",
    81  		)
    82  		pw.testRunner = "behave.py"
    83  	} else {
    84  		pw.testIncludes = append(commonIncludes,
    85  			".bootstrap/xmlrunner")
    86  
    87  		pw.testRunner = "unittest.py"
    88  	}
    89  }
    90  
    91  // Write writes the pex to the given output file.
    92  func (pw *Writer) Write(out, moduleDir string) error {
    93  	f := zip.NewFile(out, true)
    94  	defer f.Close()
    95  
    96  	// Write preamble (i.e. the shebang that makes it executable)
    97  	if err := f.WritePreamble([]byte(pw.shebang)); err != nil {
    98  		return err
    99  	}
   100  
   101  	// Write required pex stuff for tests. Note that this executable is also a zipfile and we can
   102  	// jarcat it directly in (nifty, huh?).
   103  	if len(pw.testSrcs) != 0 {
   104  		f.Include = pw.testIncludes
   105  		if err := f.AddZipFile(os.Args[0]); err != nil {
   106  			return err
   107  		}
   108  	}
   109  
   110  	// Always write pex_main.py, with some templating.
   111  	b := MustAsset("pex_main.py")
   112  	b = bytes.Replace(b, []byte("__MODULE_DIR__"), []byte(strings.Replace(moduleDir, ".", "/", -1)), 1)
   113  	b = bytes.Replace(b, []byte("__ENTRY_POINT__"), []byte(pw.realEntryPoint), 1)
   114  	b = bytes.Replace(b, []byte("__ZIP_SAFE__"), []byte(pythonBool(pw.zipSafe)), 1)
   115  
   116  	if len(pw.testSrcs) != 0 {
   117  		// If we're writing a test, we append test_main.py to it.
   118  		b2 := MustAsset("test_main.py")
   119  		b2 = bytes.Replace(b2, []byte("__TEST_NAMES__"), []byte(strings.Join(pw.testSrcs, ",")), 1)
   120  		b = append(b, b2...)
   121  		// It also needs an appropriate test runner.
   122  		b = append(b, MustAsset(pw.testRunner)...)
   123  	}
   124  	// We always append the final if __name__ == '__main__' bit.
   125  	b = append(b, MustAsset("pex_run.py")...)
   126  	return f.WriteFile("__main__.py", b, 0644)
   127  }
   128  
   129  // pythonBool returns a Python bool representation of a Go bool.
   130  func pythonBool(b bool) string {
   131  	if b {
   132  		return "True"
   133  	}
   134  	return "False"
   135  }
   136  
   137  // toPythonPath converts a normal path to a Python import path.
   138  func toPythonPath(p string) string {
   139  	ext := path.Ext(p)
   140  	return strings.Replace(p[:len(p)-len(ext)], "/", ".", -1)
   141  }