# Copyright (C) 2015-2016   Niklas Rosenstein
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import copy, glob, os, pkgutil, sys, traceback, zipfile

if sys.version_info[0] == 2:
  # FIXME: pyminifier introduces an additional, badly indented 'pass'
  # statement to single line functions. See pyminifier#89
  def items(x):
    return x.items()
  def iteritems(x):
    return x.iteritems()
else:
  def items(x):
    return list(x.items())
  def iteritems(x):
    return x.items()

class localimport(object):
  ''' Secure import mechanism that restores the previous global importer
  state after the context-manager exits. Modules imported from the local
  site will be moved into :attr:`localimport.modules`.
  Features:
    - Supports namespace packages
    - Supports evaluating ``.pth`` files
    - Automatically adds ``.egg`` files to the ``path`` when they are
      found in one of the directories specified on initialization.
    - Mocks :func:`pkgutil.extend_path()` to a version that also
      supports zipped Python Eggs.
    - Only moves imported modules into the shadows when they were
      imported from one of the paths specified in the ``localimport``
      object.
    - Reusable for delayed-import.
  .. code-block:: python
    with localimport('res/modules'):
      import some_package
  :param path:
    A string or a list of strings for the additional import paths
    that should be made available during the ``localimport`` context.
  :param parent_dir:
    If omitted, defaults to the parent directory of the ``__file__``
    variable from the parent stackframe. Relative paths in the *path*
    parameter will be considered relative to this directory.
  :param do_eggs:
    True by default. If True, ``.egg`` files in any of the specified
    *path* s will be added to :attr:`localimport.path` as well.
  :param do_pth:
    True by default. Evaluate ``.pth`` files found in the paths.
  .. attribute:: path
    The paths from which modules should be imported.These will also be
    used to determine if a module was imported from the local site and
    wether it should be released after the localimport is complete.
  .. attribute:: meta_path
    A list of importer objects that will be prepended to `sys.meta_path`
    during the localimport. Use of meta importers is discouraged as it
    could lead to problems determining whether a module is from the local
    site.
  .. attribute:: modules
    Dictionary of the modules imported from the local site.
  .. attribute:: in_context
    True when the localimport context-manager is active.
  .. attribute:: do_pth
    If True, ``.pth`` files are evaluated on :meth:`__enter__`.
  '''

  __author__ = 'Niklas Rosenstein <rosensteinniklas@gmail.com>'
  __version__ = '1.5.2'
  _py3k = sys.version_info[0] >= 3
  _string_types = (str,) if _py3k else (basestring,)

  def __init__(self, path, parent_dir=None, do_eggs=True, do_pth=True):
    if not parent_dir:
      frame = sys._getframe(1).f_globals
      if '__file__' in frame:
        parent_dir = os.path.dirname(os.path.abspath(frame['__file__']))

    # Convert relative paths to absolute paths with parent_dir and
    # evaluate .egg files in the specified directories.
    self.path = []
    if isinstance(path, self._string_types):
      path = [path]
    for path_name in path:
      if not os.path.isabs(path_name):
        if not parent_dir:
          raise ValueError('relative path but no parent_dir')
        path_name = os.path.join(parent_dir, path_name)
      path_name = os.path.normpath(path_name)
      self.path.append(path_name)
      if do_eggs:
        self.path.extend(glob.glob(os.path.join(path_name, '*.egg')))

    self.meta_path = []
    self.modules = {}
    self.do_pth = do_pth
    self.in_context = False

  def __enter__(self):
    # pkg_resources comes with setuptools.
    try:
      import pkg_resources
      nsdict = copy.deepcopy(pkg_resources._namespace_packages)
      declare_namespace = pkg_resources.declare_namespace
      pkg_resources.declare_namespace = self._declare_namespace
    except ImportError:
      nsdict = None
      declare_namespace = None

    # Save the global importer state.
    self.state = {
      'nsdict': nsdict,
      'declare_namespace': declare_namespace,
      'nspaths': {},
      'path': sys.path[:],
      'meta_path': sys.meta_path[:],
      'disables': {},
      'pkgutil.extend_path': pkgutil.extend_path,
    }

    # Update the systems meta path and apply function mocks.
    sys.path[:] = self.path
    sys.meta_path[:] = self.meta_path + sys.meta_path
    pkgutil.extend_path = self._extend_path

    # If this function is called not the first time, we need to
    # restore the modules that have been imported with it and
    # temporarily disable the ones that would be shadowed.
    for key, mod in items(self.modules):
      try: self.state['disables'][key] = sys.modules.pop(key)
      except KeyError: pass
      sys.modules[key] = mod

    # Evaluate .pth files.
    if self.do_pth:
      for path_name in self.path:
        for fn in glob.glob(os.path.join(path_name, '*.pth')):
          self._eval_pth(fn, path_name)

    # Add the original path to sys.path.
    sys.path += self.state['path']

    # Update the __path__ of all namespace modules.
    for key, mod in items(sys.modules):
      if mod is None:
        # Relative imports could have lead to None-entries in
        # sys.modules. Get rid of them so they can be re-evaluated.
        prefix = key.rpartition('.')[0]
        if hasattr(sys.modules.get(prefix), '__path__'):
          del sys.modules[key]
      elif hasattr(mod, '__path__'):
        self.state['nspaths'][key] = copy.copy(mod.__path__)
        mod.__path__ = pkgutil.extend_path(mod.__path__, mod.__name__)

    self.in_context = True
    return self

  def __exit__(self, *__):
    if not self.in_context:
      raise RuntimeError('context not entered')

    # Figure the difference of the original sys.path and the
    # current path. The list of paths will be used to determine
    # what modules are local and what not.
    local_paths = []
    for path in sys.path:
      if path not in self.state['path']:
        local_paths.append(path)

    # Move all meta path objects to self.meta_path that have not
    # been there before and have not been in the list before.
    for meta in sys.meta_path:
      if meta is not self and meta not in self.state['meta_path']:
        if meta not in self.meta_path:
          self.meta_path.append(meta)

    # Move all modules that shadow modules of the original system
    # state or modules that are from any of the localimport context
    # paths away.
    modules = sys.modules.copy()
    for key, mod in iteritems(modules):
      force_pop = False
      filename = getattr(mod, '__file__', None)
      if not filename and key not in sys.builtin_module_names:
        parent = key.rsplit('.', 1)[0]
        if parent in modules:
          filename = getattr(modules[parent], '__file__', None)
        else:
          force_pop = True
      if force_pop or (filename and self._is_local(filename, local_paths)):
        self.modules[key] = sys.modules.pop(key)

    # Restore the disabled modules.
    sys.modules.update(self.state['disables'])
    for key, mod in iteritems(self.state['disables']):
      try: parent_name = key.split('.')[-2]
      except IndexError: parent_name = None
      if parent_name and parent_name in sys.modules:
        parent = sys.modules[parent_name]
        setattr(parent, key.split('.')[-1], mod)

    # Restore the original __path__ value of namespace packages.
    for key, path in iteritems(self.state['nspaths']):
      try: sys.modules[key].__path__ = path
      except KeyError: pass

    # Restore the original state of the global importer.
    sys.path[:] = self.state['path']
    sys.meta_path[:] = self.state['meta_path']
    pkgutil.extend_path = self.state['pkgutil.extend_path']
    try:
      import pkg_resources
      pkg_resources.declare_namespace = self.state['declare_namespace']
      pkg_resources._namespace_packages.clear()
      pkg_resources._namespace_packages.update(self.state['nsdict'])
    except ImportError: pass

    self.in_context = False
    del self.state

  def _is_local(self, filename, pathlist):
    ''' Returns True if *filename* is a subpath of any of the paths
    contained by the `localimport`. This is used to determine if a
    module was imported from the local site. '''

    filename = os.path.abspath(filename)
    for path_name in pathlist:
      path_name = os.path.abspath(path_name)
      if self._is_subpath(filename, path_name):
        return True
    return False

  def _eval_pth(self, filename, sitedir):
    '''  Evaluates a `.pth` file (supporting `import` statements),
    and appends the result  to `sys.path`. Called from `__enter__()`. '''

    if not os.path.isfile(filename):
      return
    with open(filename, 'r') as fp:
      for index, line in enumerate(fp):
        if line.startswith('import'):
          line_fn = '{0}#{1}'.format(filename, index + 1)
          try:
            exec(compile(line, line_fn, 'exec'))
          except BaseException:
            traceback.print_exc()
        else:
          index = line.find('#')
          if index > 0: line = line[:index]
          line = line.strip()
          if not os.path.isabs(line):
            line = os.path.join(os.path.dirname(filename), line)
          line = os.path.normpath(line)
          if line and line not in sys.path:
            sys.path.insert(0, line)

  def _declare_namespace(self, package_name):
    ''' Mock for the original `pkg_resources.declare_namespace()`
    function that calls `pkgutil.extend_path()` afterwards as the
    original implementation doesn't seem to properly find all
    available namespace paths. '''

    self.state['declare_namespace'](package_name)
    mod = sys.modules[package_name]
    mod.__path__ = pkgutil.extend_path(mod.__path__, package_name)

  def _extend_path(self, pth, name):
    ''' Better implementation of `pkgutil.extend_path`  which adds
    support for zipped Python eggs. The original `pkgutil.extend_path`
    gets mocked by this function inside the localimport context. '''

    def zip_isfile(z, name):
      name.rstrip('/')
      return name in z.namelist()

    pname = os.path.join(*name.split('.'))
    zname = '/'.join(name.split('.'))
    init_py = '__init__' + os.extsep + 'py'
    init_pyc = '__init__' + os.extsep + 'pyc'
    init_pyo = '__init__' + os.extsep + 'pyo'

    mod_path = list(pth)
    for path in sys.path:
      if zipfile.is_zipfile(path):
        try:
          egg = zipfile.ZipFile(path, 'r')
          addpath = (
            zip_isfile(egg, zname + '/__init__.py') or
            zip_isfile(egg, zname + '/__init__.pyc') or
            zip_isfile(egg, zname + '/__init__.pyo'))
          fpath = os.path.join(path, path, zname)
          if addpath and fpath not in mod_path:
            mod_path.append(fpath)
        except (zipfile.BadZipfile, zipfile.LargeZipFile):
          pass  # xxx: Show a warning at least?
      else:
        path = os.path.join(path, pname)
        if os.path.isdir(path) and path not in mod_path:
          addpath = (
            os.path.isfile(os.path.join(path, init_py)) or
            os.path.isfile(os.path.join(path, init_pyc)) or
            os.path.isfile(os.path.join(path, init_pyo)))
          if addpath and path not in mod_path:
            mod_path.append(path)

    return [os.path.normpath(x) for x in mod_path]

  @staticmethod
  def _is_subpath(path, ask_dir):
    ''' Returns True if *path* points to the same or a
    subpathof *ask_dir*. '''

    try:
      relpath = os.path.relpath(path, ask_dir)
    except ValueError:
      return False  # happens on Windows if drive letters don't match
    return relpath == os.curdir or not relpath.startswith(os.pardir)

  def disable(self, module):
    ''' This method can be called to ensure that *module*
    and all its submodules are moved to a "disabled" state, thus
    being temporarily removed from the global importer state.

    If *module* is a list of strings, the function will be
    called for each of its items. '''

    if not isinstance(module, self._string_types):
      [self.disable(x) for x in module]
      return

    sub_prefix = module + '.'
    modules = {}
    for key, mod in iteritems(sys.modules):
      if key == module or key.startswith(sub_prefix):
        try: parent_name = '.'.join(key.split('.')[:-1])
        except IndexError: parent_name = None

        # Delete the child module reference from the parent module.
        modules[key] = mod
        if parent_name and parent_name in sys.modules:
          parent = sys.modules[parent_name]
          try:
            delattr(parent, key.split('.')[-1])
          except AttributeError:
            pass

    # Pop all the modules we found from sys.modules
    for key, mod in iteritems(modules):
      del sys.modules[key]
      self.state['disables'][key] = mod
#<endmin

__author__ = localimport.__author__
__version__ = localimport.__version__
exports = localimport       # for require()
_localimport = localimport  # backwards compatibility <= 1.4.15
