From f15e56f25d43672c784cb38aa92e449b03b97781 Mon Sep 17 00:00:00 2001
From: Robert Craigie <robert@craigie.dev>
Date: Wed, 15 Jan 2025 14:54:48 +0000
Subject: [PATCH] feat: bundle pyright inside wheel (#300)

---
 .github/workflows/release.yml |  6 +++-
 .github/workflows/test.yml    | 23 ++++++++++++
 pyproject.toml                |  3 +-
 scripts/download_pyright.py   | 68 +++++++++++++++++++++++++++++++++++
 setup.py                      |  2 +-
 src/pyright/_utils.py         |  6 ++++
 tests/test_main.py            |  8 ++++-
 tests/utils.py                | 14 ++++++++
 8 files changed, 126 insertions(+), 4 deletions(-)
 create mode 100644 scripts/download_pyright.py

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 57aeda2..9649065 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -17,12 +17,16 @@ jobs:
       - name: Set up Python 3.9
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: '3.11'
 
       - name: Install dependencies
         run: |
           python -m pip install -r dev-requirements.txt
 
+      - name: Download pyright dist
+        run: |
+          python scripts/download_pyright.py
+
       - name: Get version
         run: |
           echo "NEW_VERSION=$(python .github/scripts/get_version.py --compare)" >> $GITHUB_ENV
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index cc633dd..9a5c03e 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -50,6 +50,10 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r dev-requirements.txt
 
+      - name: Download pyright dist
+        run: |
+          python scripts/download_pyright.py
+
       - uses: actions/cache@v2
         with:
           path: ~/.cache/pip
@@ -124,6 +128,15 @@ jobs:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
 
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+
+      - name: Download pyright dist
+        run: |
+          python scripts/download_pyright.py
+
       - name: Docker Build
         uses: docker/build-push-action@v3
         # https://github.com/docker/build-push-action/#inputs
@@ -141,6 +154,16 @@ jobs:
     runs-on: windows-latest
     steps:
       - uses: actions/checkout@v3
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+
+      - name: Download pyright dist
+        run: |
+          python scripts/download_pyright.py
+
       - name: Docker Build
         # Use --% to allow double hyphen
         # Caching not currently working since we don't use buildx yet, windows
diff --git a/pyproject.toml b/pyproject.toml
index e195cfa..b4a2479 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,7 +46,8 @@ include = [
     "tests",
 ]
 exclude = [
-    "src/pyright/_mureq.py"
+    "src/pyright/_mureq.py",
+    "src/pyright/dist",
 ]
 pythonVersion = "3.9"
 
diff --git a/scripts/download_pyright.py b/scripts/download_pyright.py
new file mode 100644
index 0000000..d2a3d85
--- /dev/null
+++ b/scripts/download_pyright.py
@@ -0,0 +1,68 @@
+import io
+import sys
+import json
+import shutil
+import tarfile
+from pathlib import Path
+
+import pyright
+from pyright import _mureq, __pyright_version__
+
+DIST_DIR = Path(pyright.__file__).parent / 'dist'
+
+
+def _should_download() -> bool:
+    if '--force' in sys.argv:
+        return True
+
+    pkg_path = DIST_DIR / 'package.json'
+    if not pkg_path.exists():
+        return True
+
+    pkg_json = json.loads(pkg_path.read_text())
+    if pkg_json['version'] == __pyright_version__:
+        print(
+            f'skipping download as the current pyright version ({__pyright_version__}) is already downloaded. use --force to override'
+        )
+        return False
+
+    return True
+
+
+def download_tarball(*, version: str) -> None:
+    if not _should_download():
+        return
+
+    if DIST_DIR.exists():
+        shutil.rmtree(DIST_DIR)
+
+    rsp = _mureq.get(f'https://registry.npmjs.org/pyright/{version}')
+    rsp.raise_for_status()
+
+    info = rsp.json()
+    tar_url = info['dist']['tarball']
+    print(f'downloading tar from {tar_url}')
+
+    rsp = _mureq.get(tar_url)
+    rsp.raise_for_status()
+
+    with tarfile.open(fileobj=io.BytesIO(rsp.body)) as tar:
+        members = tar.getmembers()
+
+        # npm tarballs will always output one `package/` directory which is
+        # not necessary for our case, so we strip out the `package/` prefix
+        for member in members:
+            if member.path.startswith('package/'):
+                member.path = member.path.replace('package/', '', 1)
+            else:
+                raise RuntimeError(f'expected tar member path to start with `package/` but got {member.path}')
+
+        tar.extractall(path=DIST_DIR, members=members)
+
+
+def main() -> None:
+    download_tarball(version=__pyright_version__)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/setup.py b/setup.py
index 6d0278f..8db5069 100755
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,7 @@
         include=['pyright', 'pyright.*'],
     ),
     package_dir={'': 'src'},
-    package_data={'': ['py.typed']},
+    package_data={'': ['py.typed', 'dist/**']},
     python_requires='>=3.7',
     include_package_data=True,
     zip_safe=False,
diff --git a/src/pyright/_utils.py b/src/pyright/_utils.py
index e43b4b7..69f4d4f 100644
--- a/src/pyright/_utils.py
+++ b/src/pyright/_utils.py
@@ -43,6 +43,12 @@ def install_pyright(args: tuple[object, ...], *, quiet: bool | None) -> Path:
                 + 'Please install the new version or set PYRIGHT_PYTHON_FORCE_VERSION to `latest`\n'
             )
 
+    if version == __pyright_version__ and env_to_bool('PYRIGHT_PYTHON_USE_BUNDLED_PYRIGHT', default=True):
+        bundled_path = Path(__file__).parent.joinpath('dist')
+        if bundled_path.exists():
+            log.debug('using bundled pyright at %s', bundled_path)
+            return bundled_path
+
     cache_dir = ROOT_CACHE_DIR / version
     cache_dir.mkdir(exist_ok=True, parents=True)
 
diff --git a/tests/test_main.py b/tests/test_main.py
index b214df1..586b6e3 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -13,8 +13,9 @@
 
 import pyright
 from pyright import __pyright_version__
-from tests.utils import assert_matches
+from tests.utils import assert_matches, is_relative_to
 from pyright.utils import maybe_decode
+from pyright._utils import install_pyright
 
 if TYPE_CHECKING:
     from _pytest.monkeypatch import MonkeyPatch
@@ -238,3 +239,8 @@ def test_package_json_in_parent_dir(tmp_path: Path, monkeypatch: MonkeyPatch) ->
         check=True,
     )
     assert proc.returncode == 0
+
+
+def test_install_pyright_uses_bundled_by_default() -> None:
+    install_path = install_pyright(tuple(), quiet=None)
+    assert is_relative_to(install_path, Path(pyright.__file__).parent)
diff --git a/tests/utils.py b/tests/utils.py
index a173e74..63ff15d 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,6 +1,11 @@
 from __future__ import annotations
 
 import re
+from os import PathLike
+from pathlib import Path
+from typing_extensions import TypeAlias
+
+StrPath: TypeAlias = 'str | PathLike[str]'
 
 
 def assert_matches(pattern: re.Pattern[str], contents: str) -> re.Match[str]:
@@ -9,3 +14,12 @@ def assert_matches(pattern: re.Pattern[str], contents: str) -> re.Match[str]:
         raise ValueError(f'Pattern, {pattern}, did not match input: {contents}')
 
     return match
+
+
+def is_relative_to(path: StrPath, to: StrPath) -> bool:
+    """Backport of Path.is_relative_to for Python < 3.9"""
+    try:
+        Path(path).relative_to(to)
+        return True
+    except ValueError:
+        return False