1"""Wheels support."""
2
3from distutils.util import get_platform
4from distutils import log
5import email
6import itertools
7import os
8import posixpath
9import re
10import zipfile
11
12import pkg_resources
13import setuptools
14from pkg_resources import parse_version
15from setuptools.extern.packaging.tags import sys_tags
16from setuptools.extern.packaging.utils import canonicalize_name
17from setuptools.extern.six import PY3
18from setuptools.command.egg_info import write_requirements
19
20
21__metaclass__ = type
22
23
24WHEEL_NAME = re.compile(
25    r"""^(?P<project_name>.+?)-(?P<version>\d.*?)
26    ((-(?P<build>\d.*?))?-(?P<py_version>.+?)-(?P<abi>.+?)-(?P<platform>.+?)
27    )\.whl$""",
28    re.VERBOSE).match
29
30NAMESPACE_PACKAGE_INIT = '''\
31try:
32    __import__('pkg_resources').declare_namespace(__name__)
33except ImportError:
34    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
35'''
36
37
38def unpack(src_dir, dst_dir):
39    '''Move everything under `src_dir` to `dst_dir`, and delete the former.'''
40    for dirpath, dirnames, filenames in os.walk(src_dir):
41        subdir = os.path.relpath(dirpath, src_dir)
42        for f in filenames:
43            src = os.path.join(dirpath, f)
44            dst = os.path.join(dst_dir, subdir, f)
45            os.renames(src, dst)
46        for n, d in reversed(list(enumerate(dirnames))):
47            src = os.path.join(dirpath, d)
48            dst = os.path.join(dst_dir, subdir, d)
49            if not os.path.exists(dst):
50                # Directory does not exist in destination,
51                # rename it and prune it from os.walk list.
52                os.renames(src, dst)
53                del dirnames[n]
54    # Cleanup.
55    for dirpath, dirnames, filenames in os.walk(src_dir, topdown=True):
56        assert not filenames
57        os.rmdir(dirpath)
58
59
60class Wheel:
61
62    def __init__(self, filename):
63        match = WHEEL_NAME(os.path.basename(filename))
64        if match is None:
65            raise ValueError('invalid wheel name: %r' % filename)
66        self.filename = filename
67        for k, v in match.groupdict().items():
68            setattr(self, k, v)
69
70    def tags(self):
71        '''List tags (py_version, abi, platform) supported by this wheel.'''
72        return itertools.product(
73            self.py_version.split('.'),
74            self.abi.split('.'),
75            self.platform.split('.'),
76        )
77
78    def is_compatible(self):
79        '''Is the wheel is compatible with the current platform?'''
80        supported_tags = set((t.interpreter, t.abi, t.platform) for t in sys_tags())
81        return next((True for t in self.tags() if t in supported_tags), False)
82
83    def egg_name(self):
84        return pkg_resources.Distribution(
85            project_name=self.project_name, version=self.version,
86            platform=(None if self.platform == 'any' else get_platform()),
87        ).egg_name() + '.egg'
88
89    def get_dist_info(self, zf):
90        # find the correct name of the .dist-info dir in the wheel file
91        for member in zf.namelist():
92            dirname = posixpath.dirname(member)
93            if (dirname.endswith('.dist-info') and
94                    canonicalize_name(dirname).startswith(
95                        canonicalize_name(self.project_name))):
96                return dirname
97        raise ValueError("unsupported wheel format. .dist-info not found")
98
99    def install_as_egg(self, destination_eggdir):
100        '''Install wheel as an egg directory.'''
101        with zipfile.ZipFile(self.filename) as zf:
102            self._install_as_egg(destination_eggdir, zf)
103
104    def _install_as_egg(self, destination_eggdir, zf):
105        dist_basename = '%s-%s' % (self.project_name, self.version)
106        dist_info = self.get_dist_info(zf)
107        dist_data = '%s.data' % dist_basename
108        egg_info = os.path.join(destination_eggdir, 'EGG-INFO')
109
110        self._convert_metadata(zf, destination_eggdir, dist_info, egg_info)
111        self._move_data_entries(destination_eggdir, dist_data)
112        self._fix_namespace_packages(egg_info, destination_eggdir)
113
114    @staticmethod
115    def _convert_metadata(zf, destination_eggdir, dist_info, egg_info):
116        def get_metadata(name):
117            with zf.open(posixpath.join(dist_info, name)) as fp:
118                value = fp.read().decode('utf-8') if PY3 else fp.read()
119                return email.parser.Parser().parsestr(value)
120
121        wheel_metadata = get_metadata('WHEEL')
122        # Check wheel format version is supported.
123        wheel_version = parse_version(wheel_metadata.get('Wheel-Version'))
124        wheel_v1 = (
125            parse_version('1.0') <= wheel_version < parse_version('2.0dev0')
126        )
127        if not wheel_v1:
128            raise ValueError(
129                'unsupported wheel format version: %s' % wheel_version)
130        # Extract to target directory.
131        os.mkdir(destination_eggdir)
132        zf.extractall(destination_eggdir)
133        # Convert metadata.
134        dist_info = os.path.join(destination_eggdir, dist_info)
135        dist = pkg_resources.Distribution.from_location(
136            destination_eggdir, dist_info,
137            metadata=pkg_resources.PathMetadata(destination_eggdir, dist_info),
138        )
139
140        # Note: Evaluate and strip markers now,
141        # as it's difficult to convert back from the syntax:
142        # foobar; "linux" in sys_platform and extra == 'test'
143        def raw_req(req):
144            req.marker = None
145            return str(req)
146        install_requires = list(sorted(map(raw_req, dist.requires())))
147        extras_require = {
148            extra: sorted(
149                req
150                for req in map(raw_req, dist.requires((extra,)))
151                if req not in install_requires
152            )
153            for extra in dist.extras
154        }
155        os.rename(dist_info, egg_info)
156        os.rename(
157            os.path.join(egg_info, 'METADATA'),
158            os.path.join(egg_info, 'PKG-INFO'),
159        )
160        setup_dist = setuptools.Distribution(
161            attrs=dict(
162                install_requires=install_requires,
163                extras_require=extras_require,
164            ),
165        )
166        # Temporarily disable info traces.
167        log_threshold = log._global_log.threshold
168        log.set_threshold(log.WARN)
169        try:
170            write_requirements(
171                setup_dist.get_command_obj('egg_info'),
172                None,
173                os.path.join(egg_info, 'requires.txt'),
174            )
175        finally:
176            log.set_threshold(log_threshold)
177
178    @staticmethod
179    def _move_data_entries(destination_eggdir, dist_data):
180        """Move data entries to their correct location."""
181        dist_data = os.path.join(destination_eggdir, dist_data)
182        dist_data_scripts = os.path.join(dist_data, 'scripts')
183        if os.path.exists(dist_data_scripts):
184            egg_info_scripts = os.path.join(
185                destination_eggdir, 'EGG-INFO', 'scripts')
186            os.mkdir(egg_info_scripts)
187            for entry in os.listdir(dist_data_scripts):
188                # Remove bytecode, as it's not properly handled
189                # during easy_install scripts install phase.
190                if entry.endswith('.pyc'):
191                    os.unlink(os.path.join(dist_data_scripts, entry))
192                else:
193                    os.rename(
194                        os.path.join(dist_data_scripts, entry),
195                        os.path.join(egg_info_scripts, entry),
196                    )
197            os.rmdir(dist_data_scripts)
198        for subdir in filter(os.path.exists, (
199            os.path.join(dist_data, d)
200            for d in ('data', 'headers', 'purelib', 'platlib')
201        )):
202            unpack(subdir, destination_eggdir)
203        if os.path.exists(dist_data):
204            os.rmdir(dist_data)
205
206    @staticmethod
207    def _fix_namespace_packages(egg_info, destination_eggdir):
208        namespace_packages = os.path.join(
209            egg_info, 'namespace_packages.txt')
210        if os.path.exists(namespace_packages):
211            with open(namespace_packages) as fp:
212                namespace_packages = fp.read().split()
213            for mod in namespace_packages:
214                mod_dir = os.path.join(destination_eggdir, *mod.split('.'))
215                mod_init = os.path.join(mod_dir, '__init__.py')
216                if not os.path.exists(mod_dir):
217                    os.mkdir(mod_dir)
218                if not os.path.exists(mod_init):
219                    with open(mod_init, 'w') as fp:
220                        fp.write(NAMESPACE_PACKAGE_INIT)
221