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