virtual env based pants py replacement

python_repl.py

# coding=utf-8
# Copyright 2013 Foursquare Labs Inc. All Rights Reserved.

from __future__ import absolute_import

import os
import subprocess

from pants.backend.core.tasks.task import Task
from pants.backend.python.targets.python_binary import PythonBinary
from pants.backend.python.targets.python_library import PythonLibrary
from pants.base.exceptions import TaskError

from foursquare.pants.venv import ensure_virtualenv_for_targets, scrub_env


class PythonRepl(Task):
  ''' Task to run a python REPL with the given targets available on sys.path
  '''

  @classmethod
  def register_options(cls, register):
    super(PythonRepl, cls).register_options(register)
    register('--run', action='store_true', default=False,
      help='When specified, runs the target.')
    register('--py-run-args', action='append', default=[],
      help='Args to pass for --run')
    register('--ipython', action='store_true', default=True,
      help='Use the IPython REPL instead of the standard Python REPL (default)')
    register('--py-require', action='append', default=[],
      help='Add a module requirement for the REPL.')

  @classmethod
  def supports_passthru_args(cls):
    return True

  def prepare(self, round_manager):
    super(PythonRepl, self).prepare(round_manager)
    round_manager.require_data('python')

  def execute(self):
    self.context.release_lock()
    targets = self.context.targets()

    extra_reqs = self.get_options().py_require
    if self.get_options().ipython and not self.get_options().run:
      extra_reqs.append('ipython==2.3.0')

    if self.get_options().run and len(targets) > 0:
      if not isinstance(self.context.target_roots[0], PythonBinary):
        raise Exception(str(self.context.target_roots[0].address) + ' is not a PythonBinary, so you cannot --run it.')

    source_roots = set(l.target_base for l in targets if isinstance(l, PythonLibrary))

    venv_dir = ensure_virtualenv_for_targets(targets,
                                             extra_reqs=extra_reqs,
                                             workunit_factory=self.context.new_workunit)

    if self.get_options().run and len(targets) > 0:
      source = os.path.join(self.context.target_roots[0].sources_relative_to_buildroot()[0])
      cmd = [os.path.join(venv_dir, 'bin', 'python'), source]
      cmd.extend(self.get_options().py_run_args + self.get_passthru_args())
      env = os.environ.copy()
      # add the dependencies to the pythonpath
      env['PYTHONPATH'] = os.pathsep.join(list(source_roots))
      scrub_env(env)
      self.context.log.debug('cmd=' + str(cmd) + '\n')
      self.context.log.debug('env=' + str(env) + '\n')
      result = subprocess.call(cmd, env=env)
      if result != 0:
        raise TaskError('%s returned %d' % (source, result))
    else:
      env = os.environ.copy()
      env['PYTHONPATH'] = os.pathsep.join(list(source_roots))
      scrub_env(env)

      interpreter_path = os.path.join(venv_dir, 'bin', 'python')
      if self.get_options().ipython:
        cmd = [interpreter_path, '-m', 'IPython']
      else:
        cmd = [interpreter_path, '-c', "import code; code.interact()"]

      subprocess.Popen(cmd, env=env).wait()

venv.py

# coding=utf-8
# Copyright 2014 Foursquare Labs Inc. All Rights Reserved.

from __future__ import absolute_import, print_function

import hashlib
import os
import subprocess
import sys
import traceback
import uuid

from pants.backend.python.python_requirement import PythonRequirement
from pants.backend.python.targets.python_requirement_library import PythonRequirementLibrary
from pants.base.build_environment import get_buildroot
from pants.base.workunit import WorkUnit
from pex.fetcher import Fetcher
from pex.interpreter import PythonInterpreter
from pex.resolver import resolve
from pex.translator import WheelTranslator


def venv_name(requirements):
  requirements = list(sorted(str(req.requirement) for req in requirements))
  hasher = hashlib.sha1()
  hasher.update('\n'.join(requirements))
  return hasher.hexdigest()[:10]


def std_from_workunit(workunit=None):
  if workunit:
    return workunit.output('stdout'), workunit.output('stderr')
  else:
    return sys.stdout, sys.stderr


def get_virtualenv_cmd(temp_venv_dir):
  return [
    os.path.join('dependencies', 'pvenv', 'bin', 'python'),
    '-m', 'virtualenv',
    temp_venv_dir
  ]

def get_wheel_bootstrap_cmd(venv_dir, wheel_dir):
  return [
    os.path.join(venv_dir, 'bin', 'python'),
    '-m', 'pip', 'install',
    '--no-index',
    '--find-links={wheel_dir}'.format(wheel_dir=wheel_dir),
    'wheel>=0.23.0',
  ]

def get_wheel_install_cmd(venv_dir, reqs, wheel_dir):
  return [
    os.path.join(venv_dir, 'bin', 'python'),
    '-m', 'wheel', 'install',
    '--wheel-dir', wheel_dir,
  ] + reqs


def invoke_cmd(cmd, workunit=None):
  stdout, stderr = std_from_workunit(workunit)
  try:
    print("Invoking cmd: %s" % ' '.join(cmd), file=sys.stderr)
    ret = subprocess.Popen(cmd, stdout=stdout, stderr=stderr).wait()
  except OSError as e:
    traceback.print_exc()
    raise OSError(e.errno, e.strerror + ' for cmd ' + ' '.join(cmd))
  return ret

def _ensure_venv(venv_dir, requirements, workunit_factory=None):
  wheel_dir = os.path.realpath(os.path.expanduser('~/.fs_wheels'))
  if not os.path.isdir(venv_dir):
    temp_venv_dir = venv_dir + '.' + str(hashlib.sha1(str(uuid.uuid4())).hexdigest()[:6])
    virtualenv_cmd = get_virtualenv_cmd(temp_venv_dir)
    if workunit_factory:
      with workunit_factory('virtualenv',
                            cmd=' '.join(virtualenv_cmd),
                            labels=[WorkUnit.BOOTSTRAP, WorkUnit.TOOL]) as workunit:
        invoke_cmd(virtualenv_cmd, workunit=workunit)
    else:
      invoke_cmd(virtualenv_cmd)

    wheel_bootstrap_cmd = get_wheel_bootstrap_cmd(temp_venv_dir, wheel_dir)
    if workunit_factory:
      with workunit_factory('bootstrap-venv-wheel',
                            cmd=' '.join(wheel_bootstrap_cmd),
                            labels=[WorkUnit.BOOTSTRAP, WorkUnit.TOOL]) as workunit:
        invoke_cmd(wheel_bootstrap_cmd, workunit=workunit)
    else:
      invoke_cmd(wheel_bootstrap_cmd)

    translated_packages = resolve_requirements_for_venv(temp_venv_dir, requirements)
    bootstrapped_packages = set(['setuptools', 'wheel', 'wsgiref'])
    translated_packages = [pkg for pkg in translated_packages
                           if pkg.project_name not in bootstrapped_packages]
    full_str_reqs = [str(package.as_requirement()) for package in translated_packages]

    wheel_install_cmd = get_wheel_install_cmd(temp_venv_dir, full_str_reqs, wheel_dir)
    if workunit_factory:
      with workunit_factory('wheel-install',
                            cmd=' '.join(wheel_install_cmd),
                            labels=[WorkUnit.BOOTSTRAP, WorkUnit.TOOL]) as workunit:
        ret = invoke_cmd(wheel_install_cmd, workunit=workunit)
        if ret != 0:
          stdout = workunit.output('stdout').read()
          stderr = workunit.output('stderr').read()
          print("stdout of failed wheel install:\n{stdout}"
                .format(stdout=stdout),
                file=sys.stderr)
          print("stderr of failed wheel install:\n{stderr}"
                .format(stderr=stderr),
                file=sys.stderr)
    else:
        ret = invoke_cmd(wheel_install_cmd)
    assert ret == 0, "Failed to wheel install.  return code was %s" % (ret)

    try:
      os.symlink(temp_venv_dir, venv_dir)
    except OSError as e:
      print("Failed to rename!  Error was %s" % e)


def venv_supported_platform(venv_dir):
  venv_supported_tags_cmd = [
    os.path.join(venv_dir, 'bin', 'python'),
    '-c',
    'import pkg_resources; print(pkg_resources.get_supported_platform())',
  ]
  venv_supported_platform = subprocess.Popen(venv_supported_tags_cmd,
                                             stdout=subprocess.PIPE).communicate()[0]
  return venv_supported_platform.strip()


def resolve_requirements_for_venv(venv_dir, requirements):
  interpreter = PythonInterpreter.from_binary(os.path.join(venv_dir, 'bin', 'python'))
  wheel_dir = os.path.realpath(os.path.expanduser('~/.fs_wheels'))
  fetchers = [Fetcher([wheel_dir])]
  platform = venv_supported_platform(venv_dir)
  translator = WheelTranslator(platform=platform, interpreter=interpreter)
  return resolve(requirements,
                 fetchers=fetchers,
                 translator=translator,
                 interpreter=interpreter,
                 platform=platform)


def scrub_env(env):
  if 'VERSIONER_PYTHON_VERSION' in env:
    del env['VERSIONER_PYTHON_VERSION']
  if 'VERSIONER_PYTHON_PREFER_32_BIT' in env:
    del env['VERSIONER_PYTHON_PREFER_32_BIT']


def ensure_virtualenv(requirements, workunit_factory=None):
  requirements = list(sorted(requirements))
  virtualenv_base_dir = os.path.join(get_buildroot(), '.pvenvs')
  name = venv_name(requirements)
  venv_dir = os.path.join(virtualenv_base_dir, name)

  if workunit_factory:
    with workunit_factory('bootstrap-virtualenv', labels=[WorkUnit.BOOTSTRAP, WorkUnit.MULTITOOL]) as workunit:
      _ensure_venv(venv_dir, requirements, workunit_factory)
  else:
    _ensure_venv(venv_dir, requirements, workunit_factory)

  return os.path.realpath(venv_dir)


def ensure_virtualenv_for_targets(targets, extra_reqs=None, workunit_factory=None):
  extra_reqs = extra_reqs or []
  extra_reqs = [PythonRequirement(req) for req in extra_reqs if isinstance(req, basestring)]
  transitive_deps = set()
  for target in targets:
    target.walk(transitive_deps.add)
  requirement_libs = [lib for lib in transitive_deps if isinstance(lib, PythonRequirementLibrary)]
  requirements = set()
  for lib in requirement_libs:
    requirements.update(req for req in lib.payload.requirements)
  if extra_reqs:
    requirements.update(extra_reqs)
  return ensure_virtualenv(requirements, workunit_factory=workunit_factory)