Molsen/add pex root option

pex/bin/pex.py

@@ -42,9 +42,20 @@
 INVALID_ENTRY_POINT = 104
 
 
-def log(msg, v=False):
-  if v:
-    print(msg, file=sys.stderr)
+class Logger(object):
+  def _default_logger(self, msg, v):
+    if v:
+      print(msg, file=sys.stderr)
+
+  _LOGGER = _default_logger
+
+  def __call__(self, msg, v):
+    self._LOGGER(msg, v)
+
+  def set_logger(self, logger_callback):
+    self._LOGGER = logger_callback
+
+log = Logger()
 
 
 def parse_bool(option, opt_str, _, parser):
@@ -159,9 +170,9 @@ def configure_clp_pex_resolution(parser, builder):
   group.add_option(
       '--cache-dir',
       dest='cache_dir',
-      default=os.path.expanduser('~/.pex/build'),
+      default='{pex_root}/build',
       help='The local cache directory to use for speeding up requirement '
-           'lookups. [Default: %default]')
+           'lookups. [Default: ~/.pex/build]')
 
   group.add_option(
       '--cache-ttl',
@@ -266,9 +277,9 @@ def configure_clp_pex_environment(parser):
   group.add_option(
       '--interpreter-cache-dir',
       dest='interpreter_cache_dir',
-      default=os.path.expanduser('~/.pex/interpreters'),
+      default='{pex_root}/interpreters',
       help='The interpreter cache to use for keeping track of interpreter dependencies '
-           'for the pex tool. [Default: %default]')
+           'for the pex tool. [Default: ~/.pex/interpreters]')
 
   parser.add_option_group(group)
 
@@ -337,6 +348,13 @@ def configure_clp():
       callback=increment_verbosity,
       help='Turn on logging verbosity, may be specified multiple times.')
 
+  parser.add_option(
+      '--pex-root',
+      dest='pex_root',
+      default=None,
+      help='Specify the pex root used in this invocation of pex. [Default: ~/.pex]'
+  )
+
   parser.add_option(
       '--help-variables',
       action='callback',
@@ -491,18 +509,28 @@ def build_pex(args, options, resolver_option_builder):
   return pex_builder
 
 
-def main():
+def make_relative_to_root(path):
+  """Update options so that defaults are user relative to specified pex_root."""
+  return os.path.normpath(path.format(pex_root=ENV.PEX_ROOT))
+
+
+def main(args):
   parser, resolver_options_builder = configure_clp()
 
-  # split arguments early because optparse is dumb
-  args = sys.argv[1:]
   try:
     separator = args.index('--')
     args, cmdline = args[:separator], args[separator + 1:]
   except ValueError:
     args, cmdline = args, []
 
   options, reqs = parser.parse_args(args=args)
+  if options.pex_root:
+    ENV.set('PEX_ROOT', options.pex_root)
+  else:
+    options.pex_root = ENV.PEX_ROOT  # If option not specified fallback to env variable.
+
+  options.cache_dir = make_relative_to_root(options.cache_dir)
+  options.interpreter_cache_dir = make_relative_to_root(options.interpreter_cache_dir)
 
   with ENV.patch(PEX_VERBOSE=str(options.verbosity)):
     with TRACER.timed('Building pex'):
@@ -527,4 +555,5 @@ def main():
 
 
 if __name__ == '__main__':
-  main()
+  # split arguments early because optparse is dumb
+  main(sys.argv[1:])

pex/commands/bdist_pex.py

@@ -55,7 +55,7 @@ def run(self):
 
     reqs = [package_dir] + reqs
 
-    with ENV.patch(PEX_VERBOSE=str(options.verbosity)):
+    with ENV.patch(PEX_VERBOSE=str(options.verbosity), PEX_ROOT=options.pex_root):
       pex_builder = build_pex(reqs, options, options_builder)
 
     def split_and_strip(entry_point):

pex/http.py

@@ -248,7 +248,7 @@ def content(self, link):
 class CachedRequestsContext(RequestsContext):
   """A requests-based Context with CacheControl support."""
 
-  DEFAULT_CACHE = '~/.pex/cache'
+  DEFAULT_CACHE = os.path.join(ENV.PEX_ROOT, 'cache')
 
   def __init__(self, cache=None, **kw):
     self._cache = os.path.realpath(os.path.expanduser(cache or self.DEFAULT_CACHE))

pex/pex_info.py

@@ -30,7 +30,7 @@ class PexInfo(object):
   requirements: list                # list of requirements for this environment
 
   # Environment options
-  pex_root: ~/.pex                   # root of all pex-related files
+  pex_root: string                   # root of all pex-related files eg: ~/.pex
   entry_point: string                # entry point into this pex
   script: string                     # script to execute in this pex environment
                                      # at most one of script/entry_point can be specified

pex/testing.py

@@ -8,8 +8,10 @@
 import sys
 import tempfile
 import zipfile
+from collections import namedtuple
 from textwrap import dedent
 
+from .bin.pex import log, main
 from .common import safe_mkdir, safe_rmtree
 from .compatibility import nested
 from .installer import EggInstaller, Packager
@@ -173,6 +175,42 @@ def write_simple_pex(td, exe_contents, dists=None, coverage=False):
   return pb
 
 
+class IntegResults(namedtuple('results', 'output return_code exception')):
+  """Convenience object to return integration run results."""
+
+  def assert_success(self):
+    assert self.exception is None and self.return_code is None
+
+  def assert_failure(self):
+    assert self.exception or self.return_code
+
+
+def run_pex_command(args, env=None):
+  """Simulate running pex command for integration testing.
+
+  This is different from run_simple_pex in that it calls the pex command rather
+  than running a generated pex.  This is useful for testing end to end runs
+  with specific command line arguments or env options.
+  """
+  def logger_callback(_output):
+    def mock_logger(msg, v=None):
+      _output.append(msg)
+
+    return mock_logger
+
+  exception = None
+  error_code = None
+  output = []
+  log.set_logger(logger_callback(output))
+  try:
+    main(args=args)
+  except SystemExit as e:
+    error_code = e.code
+  except Exception as e:
+    exception = e
+  return IntegResults(output, error_code, exception)
+
+
 # TODO(wickman) Why not PEX.run?
 def run_simple_pex(pex, args=(), env=None):
   po = subprocess.Popen(

pex/variables.py

@@ -238,7 +238,7 @@ def PEX_ROOT(self):
     The directory location for PEX to cache any dependencies and code.  PEX must write
     not-zip-safe eggs and all wheels to disk in order to activate them.  Default: ~/.pex
     """
-    return self._get_path('PEX_ROOT', default=None)
+    return self._get_path('PEX_ROOT', default=os.path.expanduser('~/.pex'))
 
   @property
   def PEX_PATH(self):

tests/test_integration.py

@@ -8,7 +8,7 @@
 from twitter.common.contextutil import environment_as, temporary_dir
 
 from pex.compatibility import WINDOWS
-from pex.testing import run_simple_pex_test
+from pex.testing import run_pex_command, run_simple_pex_test
 from pex.util import named_temporary_file
 
 
@@ -23,6 +23,23 @@ def test_pex_raise():
   run_simple_pex_test(body, coverage=True)
 
 
+def test_pex_root():
+  with temporary_dir() as tmp_home:
+    with environment_as(HOME=tmp_home):
+      with temporary_dir() as td:
+        with temporary_dir() as output_dir:
+          env = os.environ.copy()
+          env['PEX_INTERPRETER'] = '1'
+
+          output_path = os.path.join(output_dir, 'pex.pex')
+          args = ['pex', '-o', output_path, '--not-zip-safe', '--pex-root={0}'.format(td)]
+          results = run_pex_command(args=args, env=env)
+          results.assert_success()
+          assert ['pex.pex'] == os.listdir(output_dir), 'Expected built pex file.'
+          assert [] == os.listdir(tmp_home), 'Expected empty temp home dir.'
+          assert 'build' in os.listdir(td), 'Expected build directory in tmp pex root.'
+
+
 def test_pex_interpreter():
   with named_temporary_file() as fp:
     fp.write(b"print('Hello world')")

tests/test_pex_info.py

@@ -5,9 +5,10 @@
 
 import pytest
 
+from pex.bin.pex import make_relative_to_root
 from pex.orderedset import OrderedSet
 from pex.pex_info import PexInfo
-from pex.variables import Variables
+from pex.variables import ENV, Variables
 
 
 def make_pex_info(requirements):
@@ -47,6 +48,14 @@ def test_from_empty_env():
   assert_same_info(PexInfo(info=info), PexInfo.from_env(env=environ))
 
 
+def test_make_relative():
+  with ENV.patch(PEX_ROOT='/pex_root'):
+    assert '/pex_root/interpreters' == make_relative_to_root('{pex_root}/interpreters')
+
+    #Verify the user can specify arbitrary absolute paths.
+    assert '/tmp/interpreters' == make_relative_to_root('/tmp/interpreters')
+
+
 def test_from_env():
   pex_root = os.path.realpath('/pex_root')
   environ = dict(PEX_ROOT=pex_root,