"""Download only the .py files from the latest wheel of a PyPI project.

Usage: python3 py_files.py <project> [output-dir]
Uses HTTP range requests, so only the zip index and the .py files are downloaded.
"""
import io, json, os, sys, urllib.request, zipfile


class HTTPFile(io.RawIOBase):
    """A read-only, seekable file over HTTP range requests."""

    def __init__(self, url):
        self.url, self.pos = url, 0
        req = urllib.request.Request(url, method="HEAD")
        self.size = int(urllib.request.urlopen(req).headers["Content-Length"])

    def seekable(self):
        return True

    def readable(self):
        return True

    def tell(self):
        return self.pos

    def seek(self, offset, whence=0):
        self.pos = [offset, self.pos + offset, self.size + offset][whence]
        return self.pos

    def readinto(self, buf):
        end = min(self.pos + len(buf), self.size)
        if self.pos >= end:
            return 0
        req = urllib.request.Request(self.url, headers={"Range": f"bytes={self.pos}-{end - 1}"})
        data = urllib.request.urlopen(req).read()
        buf[: len(data)] = data
        self.pos += len(data)
        return len(data)


def main(project, out="."):
    meta = json.load(urllib.request.urlopen(f"https://pypi.org/pypi/{project}/json"))
    wheels = [f for f in meta["urls"] if f["filename"].endswith(".whl")]
    if not wheels:
        sys.exit(f"{project} {meta['info']['version']} has no wheel")
    wheel = wheels[0]
    with zipfile.ZipFile(io.BufferedReader(HTTPFile(wheel["url"]), buffer_size=64 * 1024)) as zf:
        for info in zf.infolist():
            if info.filename.endswith(".py"):
                zf.extract(info, out)
                print(info.filename)
    print(f"From {wheel['filename']} ({wheel['size']:,} bytes in full)", file=sys.stderr)


if __name__ == "__main__":
    main(*sys.argv[1:3])