site_paths.py

"""
ul.paths.py

mimics the site module: process .pth files identically for zip and loose file distributions
"""
import zipfile
import sys
import os
import logging


class SiteProcessor(object):
    def __init__(self, root):
        self.root = root

    def process(self):
        for root, contents in self.pth_files():
            self.read_pth_file(root, contents)

    def pth_files(self):
        yield

    @staticmethod
    def read_pth_file(root, contents):
        for line in contents.readlines():
            if line.startswith('#'):
                continue
            if line.startswith("import"):
                exec (line.strip())
                continue
            new_path = "{0}/{1}".format(root, line.strip())
            if not new_path in sys.path:
                sys.path.append(new_path)


class FolderSiteProcessor(SiteProcessor):
    def pth_files(self):
        for roots, dirs, files in os.walk(self.root):
            for f in files:
                if f.endswith(".pth"):
                    clean_path = (roots + "/" + f).replace("\\", "/")
                    with open(clean_path, 'r') as handle:
                        yield roots, handle


class ZipSiteProcessor(SiteProcessor):
    def pth_files(self):
        archive = zipfile.ZipFile(self.root, 'r')
        all_names = archive.namelist()
        all_names.sort()
        pth_list = [p for p in all_names if p.endswith('.pth')]
        for pthfile in pth_list:
            local_pth = os.path.dirname(self.root + "/" + pthfile)
            contents = archive.open(pthfile, 'U')
            yield local_pth, contents


def include_site_files(*roots):
    """
    for every .pth file in or under each root, process the pth file
    in the same way as site.addsitedir()
    """
    logging.getLogger("bootstrap").info("path shim")
    for each_root in roots:
        if zipfile.is_zipfile(each_root):
            ZipSiteProcessor(each_root).process()
        else:
            FolderSiteProcessor(each_root).process()