#!/usr/bin/env python

#

# Copyright 2007 Google Inc.

#

# 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.

#

"""Pure Python zipfile importer.

This approximates the standard zipimport module, which isn't supported

by Google App Engine.  See PEP 302 for more information about the API

for import hooks.

Usage:

  import py_zipimport

As a side effect of importing, the module overrides sys.path_hooks,

and also creates an alias 'zipimport' for itself.  When your app is

running in Google App Engine production, you don't even need to import

it, since this is already done for you.  In the Google App Engine SDK

this module is not used; instead, the standard zipimport module is

used.

"""

__all__ = ['ZipImportError', 'zipimporter']

import os

import sys

import types

import UserDict

import zipfile

_SEARCH_ORDER = [

    ('.py', False),

    ('/__init__.py', True),

]

_zipfile_cache = {}

class ZipImportError(ImportError):

  """Exception raised by zipimporter objects."""

class zipimporter:

  """A PEP-302-style importer that can import from a zipfile.

  Just insert or append this class (not an instance) to sys.path_hooks

  and you're in business.  Instances satisfy both the 'importer' and

  'loader' APIs specified in PEP 302.

  """

  def __init__(self, path_entry):

    """Constructor.

    Args:

      path_entry: The entry in sys.path.  This should be the name of an

        existing zipfile possibly with a path separator and a prefix

        path within the archive appended, e.g. /x/django.zip or

        /x/django.zip/foo/bar.

    Raises:

      ZipImportError if the path_entry does not represent a valid

      zipfile with optional prefix.

    """

    archive = path_entry

    prefix = ''

    while not os.path.lexists(archive):

      head, tail = os.path.split(archive)

      if head == archive:

        msg = 'Nothing found for %r' % path_entry

        raise ZipImportError(msg)

      archive = head

      prefix = os.path.join(tail, prefix)

    if not os.path.isfile(archive):

      msg = 'Non-file %r found for %r' % (archive, path_entry)

      raise ZipImportError(msg)

    self.archive = archive

    self.prefix = os.path.join(prefix, '')

    self.zipfile = _zipfile_cache.get(archive)

    if self.zipfile is None:

      try:

        self.zipfile = zipfile.ZipFile(self.archive)

      except (EnvironmentError, zipfile.BadZipfile), err:

        msg = 'Can\'t open zipfile %s: %s: %s' % (self.archive,

                                                  err.__class__.__name__, err)

        import logging

        logging.warn(msg)

        raise ZipImportError(msg)

      else:

        _zipfile_cache[archive] = self.zipfile

        import logging

        logging.info('zipimporter(%r, %r)', archive, prefix)

  def __repr__(self):

    """Return a string representation matching zipimport.c."""

    name = self.archive

    if self.prefix:

      name = os.path.join(name, self.prefix)

    return '<zipimporter object "%s">' % name

  def _get_info(self, fullmodname):

    """Internal helper for find_module() and load_module().

    Args:

      fullmodname: The dot-separated full module name, e.g. 'django.core.mail'.

    Returns:

      A tuple (submodname, is_package, relpath) where:

        submodname: The final component of the module name, e.g. 'mail'.

        is_package: A bool indicating whether this is a package.

        relpath: The path to the module's source code within to the zipfile.

    Raises:

      ImportError if the module is not found in the archive.

    """

    parts = fullmodname.split('.')

    submodname = parts[-1]

    for suffix, is_package in _SEARCH_ORDER:

      relpath = os.path.join(self.prefix,

                             submodname + suffix.replace('/', os.sep))

      try:

        self.zipfile.getinfo(relpath.replace(os.sep, '/'))

      except KeyError:

        pass

      else:

        return submodname, is_package, relpath

    msg = ('Can\'t find module %s in zipfile %s with prefix %r' %

           (fullmodname, self.archive, self.prefix))

    raise ZipImportError(msg)

  def _get_source(self, fullmodname):

    """Internal helper for load_module().

    Args:

      fullmodname: The dot-separated full module name, e.g. 'django.core.mail'.

    Returns:

      A tuple (submodname, is_package, fullpath, source) where:

        submodname: The final component of the module name, e.g. 'mail'.

        is_package: A bool indicating whether this is a package.

        fullpath: The path to the module's source code including the

          zipfile's filename.

        source: The module's source code.

    Raises:

      ImportError if the module is not found in the archive.

    """

    submodname, is_package, relpath = self._get_info(fullmodname)

    fullpath = '%s%s%s' % (self.archive, os.sep, relpath)

    source = self.zipfile.read(relpath.replace(os.sep, '/'))

    source = source.replace('\r\n', '\n')

    source = source.replace('\r', '\n')

    return submodname, is_package, fullpath, source

  def find_module(self, fullmodname, path=None):

    """PEP-302-compliant find_module() method.

    Args:

      fullmodname: The dot-separated full module name, e.g. 'django.core.mail'.

      path: Optional and ignored; present for API compatibility only.

    Returns:

      None if the module isn't found in the archive; self if it is found.

    """

    try:

      submodname, is_package, relpath = self._get_info(fullmodname)

    except ImportError:

      return None

    else:

      return self

  def load_module(self, fullmodname):

    """PEP-302-compliant load_module() method.

    Args:

      fullmodname: The dot-separated full module name, e.g. 'django.core.mail'.

    Returns:

      The module object constructed from the source code.

    Raises:

      SyntaxError if the module's source code is syntactically incorrect.

      ImportError if there was a problem accessing the source code.

      Whatever else can be raised by executing the module's source code.

    """

    submodname, is_package, fullpath, source = self._get_source(fullmodname)

    code = compile(source, fullpath, 'exec')

    mod = sys.modules.get(fullmodname)

    try:

      if mod is None:

        mod = sys.modules[fullmodname] = types.ModuleType(fullmodname)

      mod.__loader__ = self

      mod.__file__ = fullpath

      mod.__name__ = fullmodname

      if is_package:

        mod.__path__ = [os.path.dirname(mod.__file__)]

      exec code in mod.__dict__

    except:

      if fullmodname in sys.modules:

        del sys.modules[fullmodname]

      raise

    return mod

  def get_data(self, fullpath):

    """Return (binary) content of a data file in the zipfile."""

    required_prefix = os.path.join(self.archive, '')

    if not fullpath.startswith(required_prefix):

      raise IOError('Path %r doesn\'t start with zipfile name %r' %

                    (fullpath, required_prefix))

    relpath = fullpath[len(required_prefix):]

    try:

      return self.zipfile.read(relpath)

    except KeyError:

      raise IOError('Path %r not found in zipfile %r' %

                    (relpath, self.archive))

  def is_package(self, fullmodname):

    """Return whether a module is a package."""

    submodname, is_package, relpath = self._get_info(fullmodname)

    return is_package

  def get_code(self, fullmodname):

    """Return bytecode for a module."""

    submodname, is_package, fullpath, source = self._get_source(fullmodname)

    return compile(source, fullpath, 'exec')

  def get_source(self, fullmodname):

    """Return source code for a module."""

    submodname, is_package, fullpath, source = self._get_source(fullmodname)

    return source

class ZipFileCache(UserDict.DictMixin):

  """Helper class to export archive data in _zip_directory_cache.

  Just take the info from _zipfile_cache and convert it as required.

  """

  def __init__(self, archive):

    _zipfile_cache[archive]

    self._archive = archive

  def keys(self):

    return _zipfile_cache[self._archive].namelist()

  def __getitem__(self, filename):

    info = _zipfile_cache[self._archive].getinfo(filename)

    dt = info.date_time

    dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)

    dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]

    return (os.path.join(self._archive, info.filename), info.compress_type,

            info.compress_size, info.file_size, info.header_offset, dostime,

            dosdate, info.CRC)

class ZipDirectoryCache(UserDict.DictMixin):

  """Helper class to export _zip_directory_cache."""

  def keys(self):

    return _zipfile_cache.keys()

  def __getitem__(self, archive):

    return ZipFileCache(archive)

_zip_directory_cache = ZipDirectoryCache()

sys.modules['zipimport'] = sys.modules[__name__]

sys.path_hooks[:] = [zipimporter]