A terribly janky script that downloads rust packages and makes supposedly portable dists out of them

rust-dist.py

#!/bin/python


import json
import tarfile
import shutil

from os            import environ
from sys           import exit, version_info
from typing        import Final
from pathlib       import Path

import requests
from rich          import print
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn

# Use the built-in tomllib if we support it
if version_info.minor >= 11:
    import tomllib as toml
else:
    import toml



RUST_MANIFEST_URL: Final[str] = 'https://static.rust-lang.org/dist/channel-rust-stable.toml'
MANIFEST_CACHE: Final[Path]   = (Path.cwd() / 'manifest.toml')
DL_DIR: Final[Path]           = (Path.cwd() / 'downloaded')
EXT_DIR: Final[Path]          = (Path.cwd() / 'extracted')
DIST_DIR: Final[Path]         = (Path.cwd() / 'dist')

TRIPLES: Final[tuple[str]] = (
    'x86_64-unknown-linux-gnu',
    'x86_64-unknown-linux-musl',
    'x86_64-pc-windows-msvc',
    'x86_64-pc-windows-gnu',
    'x86_64-unknown-freebsd',
    'x86_64-apple-darwin',
    'aarch64-unknown-linux-musl',
    'aarch64-pc-windows-msvc',
    'aarch64-apple-darwin',
)

PACKAGES: Final[tuple[str]] = (
    'rustc',
    'rust-std',
    'rust-mingw',
    'cargo',
)

def get_manifest(refresh: bool = False) -> dict:

    if not MANIFEST_CACHE.exists() or refresh == True:
        print('Manifest not cached, or refreshing cached version')
        r = requests.get(RUST_MANIFEST_URL)
        with MANIFEST_CACHE.open('w') as f:
            f.write(r.text)

        mnfst = toml.loads(r.text)

        print('Dumping to JSON')

        with (Path.cwd() / 'manifest.json').open('w') as f:
            json.dump(mnfst, f)

        return mnfst

    print('Using cached manifest')
    with MANIFEST_CACHE.open('rb') as f:
        return toml.load(f)

def download(url: str, target: Path, p: Progress, overwrite: bool = False):
    if target.exists() and not overwrite:
        print(f'Already have {target.name}, skipping')
        return

    with requests.get(url, stream = True) as r, target.open('wb') as f:
        r.raise_for_status()
        dl = p.add_task(f' ==> DL: {target.name}', total = int(r.headers['content-length']))
        for chunk in r.iter_content(chunk_size = 8192):
            f.write(chunk)
            p.update(dl, advance = 8192)

    p.remove_task(dl)


def extract(file: Path, dir: Path, p: Progress):
    if dir.exists():
        print(f'Already extracted {file.name}, skipping')
        return

    with file.open('rb') as f, tarfile.open(mode = 'r:xz', fileobj = f) as tar:
        ex = p.add_task(f' ==> EX: {file.name}', total = len(tar.getnames()))
        for m in tar:
            if m.isfile():
                tgt = (dir / Path(*Path(m.name).parts[1:]))
                d = tgt.parent
                if not d.exists():
                    d.mkdir(parents = True)
                with tgt.open('wb') as f:
                    with tar.extractfile(m) as tf:
                        f.write(tf.read())
                tgt.chmod(m.mode)
                p.update(ex, advance = 1)

    p.remove_task(ex)

def mkdist(ext_dir: Path, pkg: str, triple: str, p: Progress):
    dist = (DIST_DIR / triple)
    pkg_dir = (ext_dir / pkg) / pkg

    if not dist.exists():
        dist.mkdir()
    cpy = p.add_task(' ==> COPY', total = 4)
    for sd in ('bin', 'lib', 'libexec', 'share'):
        dir = (pkg_dir / sd)
        if dir.exists():
            shutil.copytree(dir, (dist / sd), dirs_exist_ok = True)
        p.update(cpy, description = f' ==> COPY {dir.name}', advance = 1)

    p.remove_task(cpy)

def main() -> int:

    if not DL_DIR.exists():
        DL_DIR.mkdir()

    if not DIST_DIR.exists():
        DIST_DIR.mkdir()

    manifest = get_manifest()
    print(f'Got rust manifest version {manifest["manifest-version"]} generated on {manifest["date"]}')

    with Progress(
        SpinnerColumn(),
        TextColumn('[progress.description]{task.description}'),
        BarColumn(bar_width = None),
        TextColumn(' {task.completed}/{task.total}'),
        transient = True
    ) as progress:
        trip = progress.add_task('', total = len(TRIPLES))

        for triple in TRIPLES:
            progress.update(trip, description = f'{triple}')

            ext     = (EXT_DIR / triple)
            pkgfile = (DIST_DIR / f'rust-{triple}.tar.xz')
            dist    = (DIST_DIR / triple)

            if not ext.exists():
                ext.mkdir(parents = True)

            paks = progress.add_task('', total = len(PACKAGES))

            for pkg in PACKAGES:
                progress.update(paks, description = f' => {pkg}')

                filename = DL_DIR / f'{triple}-{pkg}.tar.xz'
                package  = manifest['pkg'][pkg]['target']

                if triple not in package.keys():
                    continue

                if package[triple]['available']:
                    url = package[triple]['xz_url']
                    download(url, filename, progress)
                    extract(filename, ext / pkg, progress)
                    mkdist(ext, pkg, triple, progress)
                else:
                    print(f'Package {filename.stem} is not available')

                progress.update(paks, advance = 1)

            progress.remove_task(paks)

            progress.update(trip, description = f'{triple}')


            if pkgfile.exists():
                print(f'Already found dist {pkgfile.name}, skipping')
                progress.update(trip, advance = 1)
                continue

            di = progress.add_task(f' => DIST:  {pkgfile.name}', total = 4)
            with pkgfile.open('wb') as f, tarfile.open(mode = 'w:xz', fileobj = f) as tar:
                for sd in ('bin', 'lib', 'libexec', 'share'):
                    tar.add((dist / sd))
                    progress.update(di, advance = 1)

            progress.remove_task(di)

            progress.update(trip, advance = 1)
        progress.remove_task(trip)

    return 0


if __name__ == '__main__':
    exit(main())