
__author__ = 'Niklas Rosenstein <rosensteinniklas@gmail.com>'
__version__ = '1.7.6'

import copy
import glob
import os
import pkgutil
import sys
import traceback
import typing as t
import zipfile

if t.TYPE_CHECKING:
  from sys import _MetaPathFinder


def is_local(filename: str, pathlist: t.List[str]) -> bool:
  ''' Returns True if *filename* is a subpath of any of the paths in *pathlist*. '''

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


def is_subpath(path: str, parent: str) -> bool:
  ''' Returns True if *path* points to the same or a subpath of *parent*. '''

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


def eval_pth(
  filename: str,
  sitedir: str,
  dest: t.Optional[t.List[str]] = None,
  imports: t.Optional[t.List[t.Tuple[str, int, str]]] = None,
) -> t.List[str]:
  ''' Evaluates a `.pth` file (including support for `import` statements), and appends the result to the list
  *dest*. If *dest* is #None, it will fall back to `sys.path`.

  If *imports* is specified, it must be a list. `import` statements will not executed but instead appended to
  that list in tuples of (*filename*, *line*, *stmt*).
  '''

  if dest is None:
    dest = sys.path

  if not os.path.isfile(filename):
    return []

  with open(filename, 'r') as fp:
    for index, line in enumerate(fp):
      if line.startswith('import'):
        if imports is None:
          exec_pth_import(filename, index+1, line)
        else:
          imports.append((filename, index+1, line))
      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 dest:
          dest.insert(0, line)

  return dest


def exec_pth_import(filename: str, lineno: int, line: str) -> None:
  line = '\n' * (lineno - 1) + line.strip()
  try:
    exec(compile(line, filename, 'exec'))
  except BaseException:
    traceback.print_exc()


def extend_path(pth: t.List[str], name: str) -> t.List[str]:
  ''' 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]


class localimport:

  def __init__(
    self,
    path: t.Union[t.List[str], str],
    parent_dir: t.Optional[str] = None,
    do_eggs: bool = True,
    do_pth: bool = True,
    do_autodisable: bool = True,
  ) -> None:
    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, str):
      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: t.List[_MetaPathFinder] = []
    self.modules: t.Dict[str, t.Any] = {}
    self.do_pth = do_pth
    self.in_context = False
    self.do_autodisable = do_autodisable
    self.pth_imports: t.List[t.Tuple[str, int, str]] = []

    if self.do_pth:
      seen = set()
      for path_name in self.path:
        for fn in glob.glob(os.path.join(path_name, '*.pth')):
          if fn in seen: continue
          seen.add(fn)
          eval_pth(fn, path_name, dest=self.path, imports=self.pth_imports)

  def __enter__(self) -> 'localimport':
    # pkg_resources comes with setuptools.
    try:
      import pkg_resources
      nsdict = copy.deepcopy(pkg_resources._namespace_packages)  # type: ignore
      declare_namespace = pkg_resources.declare_namespace
      pkg_resources.declare_namespace = self._declare_namespace  # type: ignore
    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 = extend_path  # type: ignore

    # 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 list(self.modules.items()):
      try: self.state['disables'][key] = sys.modules.pop(key)
      except KeyError: pass
      sys.modules[key] = mod

    # Evaluate imports from the .pth files, if any.
    for fn, lineno, stmt in self.pth_imports:
      exec_pth_import(fn, lineno, stmt)

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

    # Update the __path__ of all namespace modules.
    for key, mod in list(sys.modules.items()):
      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
    if self.do_autodisable:
      self.autodisable()
    return self

  def __exit__(self, *__) -> None:
    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)
    for path in self.path:
      if path not in local_paths:
        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 modules.items():
      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 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 self.state['disables'].items():
      try: parent_name = key.split('.')[-2]
      except IndexError: parent_name = None
      if parent_name and parent_name in sys.modules:
        parent_module = sys.modules[parent_name]
        setattr(parent_module, key.split('.')[-1], mod)

    # Restore the original __path__ value of namespace packages.
    for key, path_list in self.state['nspaths'].items():
      try: sys.modules[key].__path__ = path_list
      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()  # type: ignore
      pkg_resources._namespace_packages.update(self.state['nsdict'])  # type: ignore
    except ImportError: pass

    self.in_context = False
    del self.state

  def _declare_namespace(self, package_name: str) -> None:
    '''
    Mock for #pkg_resources.declare_namespace() which 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)  # type: ignore

  def discover(self) -> t.Iterable[pkgutil.ModuleInfo]:
    return pkgutil.iter_modules(self.path)

  def disable(self, module: t.Union[t.List[str], str]) -> None:
    if not isinstance(module, str):
      for module_name in module:
        self.disable(module_name)
      return

    sub_prefix = module + '.'
    modules = {}
    for key, mod in sys.modules.items():
      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 modules.items():
      del sys.modules[key]
      self.state['disables'][key] = mod

  def autodisable(self) -> None:
    for loader, name, ispkg in self.discover():
      self.disable(name)