#!/usr/bin/env python
#
# Copyright 2010 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Program that generates scripts to create virtualenvs for App Engine.

"virtualenv is a tool to create isolated Python environments."[1]
Running the App Engine SDK within such an environment is not easy
because both, virtualenv and App Engine, modify sys.path, but in
incompatible ways.

This program creates virtualenv creation scripts that deal with the
incompabilities and allows you to use virtualenv with App Engine.

Two sample environments are defined below.

You can use this module as a library and define your own
environments. Environments are defined as dictionaries, containing a
name, packages to install via 'pip', functions to include in the
bootstrap script, and finally a name for the created script. Here is
one example:

'''
CONFIG_DJANGO10 = {'name': 'Django1.0',
                   'packages': ['django==1.0.4', 'pil'],
                   'functions': ['_ae_symlink_stdlib', '_ae_patch_distutils'],
                   'script_name': 'create_django10_env.py',
                  }
'''


What do you need?
-----------------

This module has two dependencies: 'virtualenv' and 'pip'.


How do I use it?
----------------

Please note that once you created a virtual environment, you cannot
easily add additional packages to it. This is due to the quirks of
virtualenv+AppEngine. If you try anyway, remove the opcode.py symlink
before you install a new package and restore it after the installation.


Under the hood
--------------

virtualenv creates symlinks to a number of core modules and makes the
rest available through sys.path manipulation. When you run the App
Engine SDK within an virtual environment, the App Engine SDK will only
find those symlinked modules, but none of the other stdlib modules.
To deal with this, we symlink *all* modules from stdlib into the
virtual environment.

When you create bootstrap script, a 'after_install' hook is generated
and included in. This hook then triggers the symlinking needed to run
App Engine within virtualenv.


[1] http://virtualenv.openplans.org/
"""
# we follow PEP8 - pylint: disable-msg=C6409

__author__ = 'robert.schuppenies@gmail.com (Robert Schuppenies)'

import inspect
import optparse
import os
import stat
import sys
import textwrap

import virtualenv


USAGE = """%prog ACTION [PARAMS]

Create a virtualenv bootstrap script for App Engine apps.

Actions must be one of:
  list: Prints a list of defined environments.
  create ENV_NAME: Create script for ENV_NAME."""


# Required prefix for dictionaries that contain an environment
# description.
CONFIG_PREFIX = 'CONFIG_'

# A virtual environment that has Django 1.0 installed.
CONFIG_DJANGO10 = {'name': 'Django1.0',
                   'packages': ['django==1.0.4'],
                   'functions': ['_ae_symlink_stdlib', '_ae_patch_distutils'],
                   'script_name': 'create_django10_env.py',
                  }

# A virtual environment that has Django 1.0 installed.
CONFIG_DJANGO11 = {'name': 'Django1.1',
                   'packages': ['django==1.1.2'],
                   'functions': ['_ae_symlink_stdlib', '_ae_patch_distutils'],
                   'script_name': 'create_django11_env.py',
                  }


def _ae_symlink_stdlib(home_dir):
  """Create symlinks of stdlib modules of the current interpreter.

  This function will be included in the virtualenv bootstrap script.

  Args:
    home_dir: Name of the virtual environment.
  """
  # This function is included in the bootstrap script, so we have to
  # do the import inside the function - pylint: disable-msg=C6204
  import os
  # 'path_locations' is defined in the bootstrap script. We can assume
  # it to be there - pylint: disable-msg=E0602
  home_dir, lib_dir, _, _ = path_locations(home_dir)
  # Use 'os' module because it should always exist.
  source_dir = os.path.dirname(sys.modules['os'].__file__)
  for f in os.listdir(source_dir):
    if f.endswith('.py'):
      full_source = os.path.join(source_dir, f)
      full_target = os.path.join(lib_dir, f)
      if not os.path.exists(full_target):
        os.symlink(full_source, full_target)


def _ae_patch_distutils(home_dir):
  """Patch distutils wrapper script.

  The distutils wrapper script used by virtual_env has a hard-coded
  dependency to the location of the actual distutils module. This
  location is incorrect for our case so we need to patch it.

  This function will be included in the virtualenv bootstrap script.

  Args:
    home_dir: Name of the virtual environment.
  """
  # This function is included in the bootstrap script, so we have to
  # do the import inside the function - pylint: disable-msg=C6204
  import fileinput, os
  # String in virtualenv distutil wrapper module used to locate the
  # original distutil module.
  DIST_UTIL_PATH_STR = ('distutils_path = os.path.join('
                        'os.path.dirname(opcode.__file__), \'distutils\')')
  # Use 'os' module because it does always exist.
  source_dir = os.path.dirname(sys.modules['os'].__file__)
  # 'path_locations' is defined in the bootstrap script. We can assume
  # it to be there - pylint: disable-msg=E0602
  home_dir, lib_dir, _, _ = path_locations(home_dir)
  distutils_init = os.path.join(lib_dir, 'distutils', '__init__.py')
  replace = 'distutils_path = "%s"' % os.path.join(source_dir, 'distutils')
  for line in fileinput.FileInput(distutils_init, inplace=1):
    line = line.replace(DIST_UTIL_PATH_STR, replace).replace(os.linesep, '')
    print line


def generate_after_install_source(packages=None, functions=None):
  """Generate the source code for 'after_install'.

  Custom bootstrap scripts for virtualenv can define an
  'after_install' function which is invoked once a virtual environment
  has been created. The 'after_install' function can then be used to
  trigger additional customizations of the new environment, such as
  the installation of default packages. We use this hook to invoke
  code that makes virtualenv and App Engine work together.

  The package names you specify in 'packages' will be installed into
  the new virtual environment with pip.

  Functions defined in 'functions' are added as separate functions to
  the generated source code. Be aware that they are inserted into a
  predefined bootstrap script of virtualenv and thus name collisions
  are possible.

  Args:
    packages: A list of packages that should be installed via pip.
    functions: A list of functions to be included in the source.

  Returns:
    The source code for 'after_install'.
  """
  if not packages:
    packages = []
  if not functions:
    functions = []
  source = []
  source += [inspect.getsource(function) for function in functions]
  source.append('def after_install(options, home_dir):')
  for package in packages:
    cmd = ("  subprocess.call(['pip', 'install', '-E', home_dir, '%s'])" %
           package)
    source.append(cmd)
  source.append('  _ae_symlink_stdlib(home_dir)')
  source.append('  _ae_patch_distutils(home_dir)')
  return os.linesep.join(source)


def create_bootstrap_script(config):
  """Create a bootstrap script based on 'config'.

  'config' should be a dictionary with the following key value pairs:
    {'name': 'CONFIGURATION_NAME',
     'packages': LIST_OF_PACKAGE_NAMES,        # required
     'functions': LIST_OF_FUNCTION_NAMES,      # required
     'script_name': NAME_FOR_GENERATED_SCRIPT, # required
    }


  The created script is stored to config['script_name'].

  Args:
    config: An environment configuration.
  """
  source = generate_after_install_source(config['packages'],
                                         config['functions'])
  source = textwrap.dedent(source)
  output = virtualenv.create_bootstrap_script(source)
  script_name = config['script_name']
  f = open(script_name, 'w')
  f.write(output)
  f.close()
  os.chmod(script_name, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)


def _get_env_configs():
  """Get a list of all predefined environment configurations.

  Returns:
    A dictionary of (config name, config object) entries.
  """
  configs = {}
  this_module = inspect.getmodule(lambda: None)
  for arg in dir(this_module):
    if arg.startswith(CONFIG_PREFIX):
      obj = getattr(this_module, arg)
      if isinstance(obj, dict):
        config = dict(obj)
        functions = []
        for function_name in obj['functions']:
          func = getattr(this_module, function_name)
          functions.append(func)
        config['functions'] = functions
        configs[obj['name']] = config
  return configs


def _action_list_environments():
  """Print all defined environment configurations."""
  configs = _get_env_configs()
  for name, config in configs.items():
    print '"%s"' % name
    print '-'*(len(name) + 2)
    print 'packages: %s' % ', '.join(config['packages'])
    print 'script name: %s' % config['script_name']
    print


def _action_create_bootstrap_script(config_name):
  """Create a bootstrap script for 'config_name'.

  Args:
    config_name: Name of configuration for which a script will be created.
  """
  configs = _get_env_configs()
  if config_name not in configs:
    print ('Error: "%s" does not exist. I only know those: %s' %
           (config_name, configs.keys()))
    sys.exit(1)
  config = configs[config_name]
  create_bootstrap_script(config)
  script_name = config['script_name']
  print 'Bootstrap script generated.You can create a virtual environment with:'
  print '$ ./%s --distribute --no-site-packages ENV_NAME' % script_name


if __name__ == '__main__':
  parser = optparse.OptionParser(USAGE)
  (options, args) = parser.parse_args()
  if len(args) <= 0:
    print 'Error: No action defined.'
    parser.print_help()
    sys.exit(1)
  ACTION = args[0]
  if ACTION == 'list':
    _action_list_environments()
  elif ACTION == 'create':
    if len(args) != 2:
      print 'Error: Expected a single name.'
      sys.exit(1)
    _action_create_bootstrap_script(args[1])
  else:
    print 'Error: Unknown action "%s"' % ACTION
    sys.exit(1)
