From 46a070fb7f877054a5db0c19ab29d0c3cf2eebdd Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Fri, 5 Feb 2021 23:39:07 -0600
Subject: [PATCH] Add whlimport wasm module

---
 wasm/demo/snippets/import_pypi.py |  22 +++++++
 wasm/lib/Lib/whlimport.py         | 102 ++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+)
 create mode 100644 wasm/demo/snippets/import_pypi.py
 create mode 100644 wasm/lib/Lib/whlimport.py

diff --git a/wasm/demo/snippets/import_pypi.py b/wasm/demo/snippets/import_pypi.py
new file mode 100644
index 000000000..9acf3d513
--- /dev/null
+++ b/wasm/demo/snippets/import_pypi.py
@@ -0,0 +1,22 @@
+import asyncweb
+import pypimport
+
+pypimport.setup()
+
+# shim path utilities into the "os" module
+class os:
+    import posixpath as path
+import sys
+sys.modules['os'] = os
+sys.modules['os.path'] = os.path
+del sys, os
+
+@asyncweb.main
+async def main():
+    await pypimport.load_package("pygments")
+    import pygments
+    import pygments.lexers
+    import pygments.formatters.html
+    lexer = pygments.lexers.get_lexer_by_name("python")
+    fmter = pygments.formatters.html.HtmlFormatter(noclasses=True, style="default")
+    print(pygments.highlight("print('hi, mom!')", lexer, fmter))
diff --git a/wasm/lib/Lib/whlimport.py b/wasm/lib/Lib/whlimport.py
new file mode 100644
index 000000000..e58fb733d
--- /dev/null
+++ b/wasm/lib/Lib/whlimport.py
@@ -0,0 +1,102 @@
+import browser
+import zipfile
+import asyncweb
+import io
+import _frozen_importlib as _bootstrap
+
+_IS_SETUP = False
+def setup(*, log=print):
+    global _IS_SETUP, LOG_FUNC
+
+    if not _IS_SETUP:
+        import sys
+        sys.meta_path.insert(0, WheelFinder)
+        _IS_SETUP = True
+
+    if not log:
+        def LOG_FUNC(log):
+            pass
+    else:
+        LOG_FUNC = log
+
+async def load_package(*args):
+    await asyncweb.wait_all(_load_package(pkg) for pkg in args)
+
+_loaded_packages = {}
+
+LOG_FUNC = print
+
+async def _load_package(pkg):
+    # TODO: support pkg==X.Y semver specifiers as well as arbitrary URLs
+    info = await browser.fetch(f'https://pypi.org/pypi/{pkg}/json', response_format="json")
+    name = info['info']['name']
+    ver = info['info']['version']
+    ver_downloads = info['releases'][ver]
+    try:
+        dl = next(dl for dl in ver_downloads if dl['packagetype'] == 'bdist_wheel')
+    except StopIteration:
+        raise ValueError(f"no wheel available for package {Name!r} {ver}")
+    if name in _loaded_packages:
+        return
+    fname = dl['filename']
+    LOG_FUNC(f"Downloading {fname} ({format_size(dl['size'])})...")
+    zip_data = io.BytesIO(await browser.fetch(dl['url'], response_format="array_buffer"))
+    size = len(zip_data.getbuffer())
+    LOG_FUNC(f"{fname} done!")
+    _loaded_packages[name] = zipfile.ZipFile(zip_data)
+
+def format_size(bytes):
+    # type: (float) -> str
+    if bytes > 1000 * 1000:
+        return '{:.1f} MB'.format(bytes / 1000.0 / 1000)
+    elif bytes > 10 * 1000:
+        return '{} kB'.format(int(bytes / 1000))
+    elif bytes > 1000:
+        return '{:.1f} kB'.format(bytes / 1000.0)
+    else:
+        return '{} bytes'.format(int(bytes))
+
+class WheelFinder:
+    _packages = _loaded_packages
+    
+    @classmethod
+    def find_spec(cls, fullname, path=None, target=None):
+        path = fullname.replace('.', '/')
+        for zname, z in cls._packages.items():
+            mi, fullpath = _get_module_info(z, path)
+            if mi is not None:
+                return _bootstrap.spec_from_loader(fullname, cls, origin=f'wheel:{zname}/{fullpath}', is_package=mi)
+        return None
+    
+    @classmethod
+    def create_module(cls, spec):
+        return None
+    
+    @classmethod
+    def exec_module(cls, module):
+        origin = module.__spec__.origin
+        if not origin or not origin.startswith("wheel:"):
+            raise ImportError(f'{module.__spec__.name!r} is not a zip module')
+
+        zipname, slash, path = origin[len('wheel:'):].partition('/')
+        source = cls._packages[zipname].read(path)
+        code = _bootstrap._call_with_frames_removed(compile, source, origin, 'exec', dont_inherit=True)
+        _bootstrap._call_with_frames_removed(exec, code, module.__dict__)
+
+
+_zip_searchorder = (
+    # (path_sep + '__init__.pyc', True, True),
+    ('/__init__.py', False, True),
+    # ('.pyc', True, False),
+    ('.py', False, False),
+)
+
+def _get_module_info(zf, path):
+    for suffix, isbytecode, ispackage in _zip_searchorder:
+        fullpath = path + suffix
+        try:
+            zf.getinfo(fullpath)
+        except KeyError:
+            continue
+        return ispackage, fullpath
+    return None, None