|
| 1 | +import browser |
| 2 | +import zipfile |
| 3 | +import asyncweb |
| 4 | +import io |
| 5 | +import _frozen_importlib as _bootstrap |
| 6 | + |
| 7 | +_IS_SETUP = False |
| 8 | +def setup(*, log=print): |
| 9 | + global _IS_SETUP, LOG_FUNC |
| 10 | + |
| 11 | + if not _IS_SETUP: |
| 12 | + import sys |
| 13 | + sys.meta_path.insert(0, WheelFinder) |
| 14 | + _IS_SETUP = True |
| 15 | + |
| 16 | + if not log: |
| 17 | + def LOG_FUNC(log): |
| 18 | + pass |
| 19 | + else: |
| 20 | + LOG_FUNC = log |
| 21 | + |
| 22 | +async def load_package(*args): |
| 23 | + await asyncweb.wait_all(_load_package(pkg) for pkg in args) |
| 24 | + |
| 25 | +_loaded_packages = {} |
| 26 | + |
| 27 | +LOG_FUNC = print |
| 28 | + |
| 29 | +async def _load_package(pkg): |
| 30 | + # TODO: support pkg==X.Y semver specifiers as well as arbitrary URLs |
| 31 | + info = await browser.fetch(f'https://pypi.org/pypi/{pkg}/json', response_format="json") |
| 32 | + name = info['info']['name'] |
| 33 | + ver = info['info']['version'] |
| 34 | + ver_downloads = info['releases'][ver] |
| 35 | + try: |
| 36 | + dl = next(dl for dl in ver_downloads if dl['packagetype'] == 'bdist_wheel') |
| 37 | + except StopIteration: |
| 38 | + raise ValueError(f"no wheel available for package {Name!r} {ver}") |
| 39 | + if name in _loaded_packages: |
| 40 | + return |
| 41 | + fname = dl['filename'] |
| 42 | + LOG_FUNC(f"Downloading {fname} ({format_size(dl['size'])})...") |
| 43 | + zip_data = io.BytesIO(await browser.fetch(dl['url'], response_format="array_buffer")) |
| 44 | + size = len(zip_data.getbuffer()) |
| 45 | + LOG_FUNC(f"{fname} done!") |
| 46 | + _loaded_packages[name] = zipfile.ZipFile(zip_data) |
| 47 | + |
| 48 | +def format_size(bytes): |
| 49 | + # type: (float) -> str |
| 50 | + if bytes > 1000 * 1000: |
| 51 | + return '{:.1f} MB'.format(bytes / 1000.0 / 1000) |
| 52 | + elif bytes > 10 * 1000: |
| 53 | + return '{} kB'.format(int(bytes / 1000)) |
| 54 | + elif bytes > 1000: |
| 55 | + return '{:.1f} kB'.format(bytes / 1000.0) |
| 56 | + else: |
| 57 | + return '{} bytes'.format(int(bytes)) |
| 58 | + |
| 59 | +class WheelFinder: |
| 60 | + _packages = _loaded_packages |
| 61 | + |
| 62 | + @classmethod |
| 63 | + def find_spec(cls, fullname, path=None, target=None): |
| 64 | + path = fullname.replace('.', '/') |
| 65 | + for zname, z in cls._packages.items(): |
| 66 | + mi, fullpath = _get_module_info(z, path) |
| 67 | + if mi is not None: |
| 68 | + return _bootstrap.spec_from_loader(fullname, cls, origin=f'wheel:{zname}/{fullpath}', is_package=mi) |
| 69 | + return None |
| 70 | + |
| 71 | + @classmethod |
| 72 | + def create_module(cls, spec): |
| 73 | + return None |
| 74 | + |
| 75 | + @classmethod |
| 76 | + def exec_module(cls, module): |
| 77 | + origin = module.__spec__.origin |
| 78 | + if not origin or not origin.startswith("wheel:"): |
| 79 | + raise ImportError(f'{module.__spec__.name!r} is not a zip module') |
| 80 | + |
| 81 | + zipname, slash, path = origin[len('wheel:'):].partition('/') |
| 82 | + source = cls._packages[zipname].read(path) |
| 83 | + code = _bootstrap._call_with_frames_removed(compile, source, origin, 'exec', dont_inherit=True) |
| 84 | + _bootstrap._call_with_frames_removed(exec, code, module.__dict__) |
| 85 | + |
| 86 | + |
| 87 | +_zip_searchorder = ( |
| 88 | + # (path_sep + '__init__.pyc', True, True), |
| 89 | + ('/__init__.py', False, True), |
| 90 | + # ('.pyc', True, False), |
| 91 | + ('.py', False, False), |
| 92 | +) |
| 93 | + |
| 94 | +def _get_module_info(zf, path): |
| 95 | + for suffix, isbytecode, ispackage in _zip_searchorder: |
| 96 | + fullpath = path + suffix |
| 97 | + try: |
| 98 | + zf.getinfo(fullpath) |
| 99 | + except KeyError: |
| 100 | + continue |
| 101 | + return ispackage, fullpath |
| 102 | + return None, None |
0 commit comments