github.py

from __future__ import annotations

import json
import os
import random
import sys
import time
import urllib.request
import warnings
from collections.abc import Iterable, Iterator, Mapping, Sequence
from itertools import islice
from pathlib import Path
from typing import IO, TYPE_CHECKING, Any, ClassVar, Literal, TypeVar, cast

import polars as pl
from polars import col

from tools.datasets import semver
from tools.datasets.models import (
    GitHubRateLimitResources,
    GitHubTag,
    GitHubTree,
    GitHubTreesResponse,
    GitHubUrl,
    ParsedRateLimit,
    ParsedTag,
    ParsedTree,
    SemVerTag,
)

if sys.version_info >= (3, 13):
    from typing import is_typeddict
else:
    from typing_extensions import is_typeddict

if TYPE_CHECKING:
    from collections.abc import MutableMapping
    from email.message import Message
    from urllib.request import OpenerDirector, Request

    from altair.datasets._typing import Extension

    if sys.version_info >= (3, 13):
        from typing import TypeIs
    else:
        from typing_extensions import TypeIs
    if sys.version_info >= (3, 11):
        from typing import LiteralString
    else:
        from typing_extensions import LiteralString
    if sys.version_info >= (3, 10):
        from typing import TypeAlias
    else:
        from typing_extensions import TypeAlias

    _PathName: TypeAlias = Literal["dir", "tags", "trees"]


__all__ = ["GitHub"]

_TD = TypeVar("_TD", bound=Mapping[str, Any])

_DATA = "data"


def is_ext_supported(suffix: str) -> TypeIs[Extension]:
    return suffix in {".csv", ".json", ".tsv", ".arrow", ".parquet"}


def _is_str(obj: Any) -> TypeIs[str]:
    return isinstance(obj, str)


class _ErrorHandler(urllib.request.BaseHandler):
    """
    Adds `rate limit`_ info to a forbidden error.

    .. _rate limit:
        https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api?apiVersion=2022-11-28
    """

    def http_error_default(
        self, req: Request, fp: IO[bytes] | None, code: int, msg: str, hdrs: Message
    ):
        if code == 403 and (reset := hdrs.get("X-RateLimit-Reset", None)):
            limit = hdrs.get("X-RateLimit-Limit", "")
            remaining = hdrs.get("X-RateLimit-Remaining", "")
            msg = (
                f"{msg}\n\nFailed to balance rate limit.\n"
                f"{limit=}, {remaining=}\n"
                f"Reset: {time.localtime(int(reset))!r}"
            )
        raise urllib.request.HTTPError(req.full_url, code, msg, hdrs, fp)


class _GitHubRequestNamespace:
    """
    Fetching resources from the `GitHub API`_.

    .. _GitHub API:
        https://docs.github.com/en/rest/about-the-rest-api/about-the-rest-api?apiVersion=2022-11-28
    """

    _ENV_VAR: LiteralString = "VEGA_GITHUB_TOKEN"
    _TAGS_MAX_PAGE: Literal[100] = 100
    _VERSION: LiteralString = "2022-11-28"
    _UNAUTH_RATE_LIMIT: Literal[60] = 60
    _TAGS_COST: Literal[1] = 1
    _TREES_COST: Literal[2] = 2
    _UNAUTH_DELAY: Literal[5_000] = 5_000
    """**ms** delay added between **unauthenticated** ``trees`` requests."""
    _AUTH_DELAY: Literal[500] = 500
    """**ms** delay added between **authenticated** ``trees`` requests."""
    _UNAUTH_TREES_LIMIT: Literal[10] = 10

    def __init__(self, gh: GitHub, /) -> None:
        self._gh = gh

    @property
    def url(self) -> GitHubUrl:
        return self._gh.url

    def rate_limit(self) -> GitHubRateLimitResources:
        with self._gh._opener.open(self._request(self.url.RATE)) as response:
            content: GitHubRateLimitResources = json.load(response)["resources"]
        return content

    def delay(self, *, is_auth: bool) -> float:
        ms = self._AUTH_DELAY if is_auth else self._UNAUTH_DELAY
        return (ms + random.triangular()) / 1_000

    def tags(self, n: int, *, warn_lower: bool) -> list[GitHubTag]:
        if n < 1 or n > self._TAGS_MAX_PAGE:
            raise ValueError(n)
        req = self._request(f"{self.url.TAGS}?per_page={n}")
        with self._gh._opener.open(req) as response:
            content: list[GitHubTag] = json.load(response)
        if warn_lower and len(content) < n:
            earliest = response[-1]["name"]
            n_response = len(content)
            msg = f"Requested {n=} tags, but got {n_response}\n{earliest=}"
            warnings.warn(msg, stacklevel=3)
        return content

    def trees(self, tag: str | ParsedTag, /) -> GitHubTreesResponse:
        """For a given ``tag``, perform **2x requests** to get directory metadata."""
        if _is_str(tag):
            url = tag if tag.startswith(self.url.TREES) else f"{self.url.TREES}{tag}"
        else:
            url = tag["trees_url"]
        with self._gh._opener.open(self._request(url)) as response:
            content: GitHubTreesResponse = json.load(response)
        query = (tree["url"] for tree in content["tree"] if tree["path"] == _DATA)
        if data_url := next(query, None):
            with self._gh._opener.open(self._request(data_url)) as response:
                data_dir: GitHubTreesResponse = json.load(response)
            return data_dir
        else:
            raise FileNotFoundError

    def _request(self, url: str, /, *, raw: bool = False) -> Request:
        """
        Wrap a request url with a `personal access token`_ - if set as an env var.

        By default the endpoint returns json, specify raw to get blob data.
        See `Media types`_.

        .. _personal access token:
            https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens
        .. _Media types:
            https://docs.github.com/en/rest/using-the-rest-api/getting-started-with-the-rest-api?apiVersion=2022-11-28#media-types
        """
        headers: MutableMapping[str, str] = {"X-GitHub-Api-Version": self._VERSION}
        if tok := os.environ.get(self._ENV_VAR):
            headers["Authorization"] = (
                tok if tok.startswith("Bearer ") else f"Bearer {tok}"
            )
        if raw:
            headers["Accept"] = "application/vnd.github.raw+json"
        return urllib.request.Request(url, headers=headers)


class _GitHubParseNamespace:
    """
    Transform responses into intermediate representations.

    Where relevant:
    - Adding cheap to compute metadata
    - Dropping information that we don't need for the task
    """

    def __init__(self, gh: GitHub, /) -> None:
        self._gh = gh

    @property
    def url(self) -> GitHubUrl:
        return self._gh.url

    def rate_limit(self, rate_limit: GitHubRateLimitResources, /) -> ParsedRateLimit:
        core = rate_limit["core"]
        reset = core["reset"]
        return ParsedRateLimit(
            **core,
            reset_time=time.localtime(reset),
            is_limited=core["remaining"] == 0,
            is_auth=core["limit"] > self._gh.req._UNAUTH_RATE_LIMIT,
        )

    def tag(self, tag: GitHubTag, /) -> ParsedTag:
        sha = tag["commit"]["sha"]
        return ParsedTag(tag=tag["name"], sha=sha, trees_url=f"{self.url.TREES}{sha}")

    def tags(self, tags: list[GitHubTag], /) -> list[ParsedTag]:
        return [self.tag(t) for t in tags]

    def tree(self, tree: GitHubTree, tag: str, /) -> ParsedTree:
        """For a single tree (file) convert to an IR with only relevant properties."""
        path = Path(tree["path"])
        return ParsedTree(
            file_name=path.name,
            dataset_name=path.stem,
            suffix=path.suffix,
            size=tree["size"],
            sha=tree["sha"],
            ext_supported=is_ext_supported(path.suffix),
            tag=tag,
        )

    def trees(self, tree: GitHubTreesResponse, /, tag: str) -> list[ParsedTree]:
        """For a tree response (directory of files) convert to an IR with only relevant properties."""
        return [self.tree(t, tag) for t in tree["tree"]]

    def tag_from_str(self, s: str, /) -> str:
        # - Actual tag
        # - Trees url (using ref name)
        # - npm url (works w/o the `v` prefix)
        trees_url = self.url.TREES
        npm_url = self._gh._npm_cdn_url
        if s.startswith("v"):
            return s
        elif s.startswith(trees_url):
            return s.replace(trees_url, "")
        elif s.startswith(npm_url):
            s, _ = s.replace(npm_url, "").split("/")
            return s if s.startswith("v") else f"v{s}"
        else:
            raise TypeError(s)


class GitHub:
    """
    Primary interface with the GitHub API.

    Maintains up-to-date metadata, describing **every** available dataset across **all known** releases.

    - Uses `tags`_, `trees`_, `rate_limit`_ endpoints.
    - Organizes distinct groups of operations into property accessor namespaces.

    .. _tags:
        https://docs.github.com/en/rest/repos/repos?apiVersion=2022-11-28#list-repository-tags
    .. _trees:
        https://docs.github.com/en/rest/git/trees?apiVersion=2022-11-28#get-a-tree
    .. _rate_limit:
        https://docs.github.com/en/rest/rate-limit/rate-limit?apiVersion=2022-11-28#get-rate-limit-status-for-the-authenticated-user
    """

    _opener: ClassVar[OpenerDirector] = urllib.request.build_opener(_ErrorHandler)

    def __init__(
        self,
        out_dir_tools: Path,
        out_dir_altair: Path,
        name_tags: str,
        name_trees: str,
        *,
        npm_cdn_url: LiteralString,
        base_url: LiteralString = "https://api.github.com/",
        org: LiteralString = "vega",
        package: LiteralString = "vega-datasets",
    ) -> None:
        out_dir_tools.mkdir(exist_ok=True)
        out_dir_altair.mkdir(exist_ok=True)
        self._paths: dict[_PathName, Path] = {
            "dir": out_dir_tools,
            "tags": out_dir_tools / f"{name_tags}.parquet",
            "trees": out_dir_altair / f"{name_trees}.parquet",
        }
        repo = f"{base_url}repos/{org}/{package}/"
        self._url = GitHubUrl(
            BASE=base_url,
            BLOBS=f"{repo}git/blobs/",
            RATE=f"{base_url}rate_limit",
            REPO=repo,
            TAGS=f"{repo}tags",
            TREES=f"{repo}git/trees/",
        )
        self._npm_cdn_url: LiteralString = npm_cdn_url

    @property
    def req(self) -> _GitHubRequestNamespace:
        return _GitHubRequestNamespace(self)

    @property
    def parse(self) -> _GitHubParseNamespace:
        return _GitHubParseNamespace(self)

    @property
    def url(self) -> GitHubUrl:
        return self._url

    def rate_limit(self, *, strict: bool = False) -> ParsedRateLimit:
        limit = self.parse.rate_limit(self.req.rate_limit())
        if strict and limit["is_limited"]:
            warnings.warn(
                f"Reached rate limit:\n{limit!r}\n\n"
                f"Try setting environment variable {self.req._ENV_VAR!r}",
                stacklevel=2,
            )
        return limit

    def delay(self, rate_limit: ParsedRateLimit | None = None, /) -> float:
        """Return a delay time in seconds, corresponding with authentication status."""
        limit = rate_limit or self.rate_limit(strict=True)
        return self.req.delay(is_auth=limit["is_auth"])

    def tags(
        self,
        n_head: int | None = None,
        *,
        npm_tags: pl.DataFrame | pl.LazyFrame | None = None,
        warn_lower: bool = False,
    ) -> pl.DataFrame:
        """
        Get release info, enhance with `SemVer`_ context.

        Parameters
        ----------
        n_head
            Limit to most recent releases.
        npm_tags
            Used to remove any github-only releases.
        warn_lower
            Emit a warning if fewer than ``n_head`` tags were returned.

        .. _SemVer:
            https://semver.org/#semantic-versioning-200
        """
        tags = self.req.tags(n_head or self.req._TAGS_MAX_PAGE, warn_lower=warn_lower)
        frame = pl.DataFrame(self.parse.tags(tags)).pipe(semver.with_columns)
        if npm_tags is not None:
            return frame.lazy().join(npm_tags.lazy().select("tag"), on="tag").collect()
        else:
            return frame

    def trees(self, tag: str | ParsedTag, /) -> pl.DataFrame:
        """Retrieve directory info for a given version ``tag``."""
        trees = self.req.trees(tag)
        tag_v = self.parse.tag_from_str(tag) if _is_str(tag) else tag["tag"]
        parsed = self.parse.trees(trees, tag=tag_v)
        url = pl.concat_str(
            pl.lit(self._npm_cdn_url),
            col("tag"),
            pl.lit(f"/{_DATA}/"),
            col("file_name"),
        )
        df = (
            pl.LazyFrame(parsed)
            .with_columns(
                name_collision=col("dataset_name").is_duplicated(), url_npm=url
            )
            .collect()
        )
        return df.select(*sorted(df.columns))

    def refresh_trees(self, gh_tags: pl.DataFrame, /) -> pl.DataFrame:
        """
        Use known tags to discover and update missing trees metadata.

        Aims to stay well-within API rate limits, both for authenticated and unauthenticated users.

        Notes
        -----
        Internally handles regenerating the ``tag`` enum.
        """
        if gh_tags.is_empty():
            msg = f"Expected rows present in `gh_tags`, but got:\n{gh_tags!r}"
            raise NotImplementedError(msg)
        rate_limit = self.rate_limit(strict=True)
        stop = None if rate_limit["is_auth"] else self.req._UNAUTH_TREES_LIMIT
        fp = self._paths["trees"]
        if not fp.exists():
            print(f"Initializing {fp!s}")
            result = self._trees_batched(_iter_rows(gh_tags, stop, SemVerTag))
        else:
            trees = (
                pl.scan_parquet(fp).with_columns(col("tag").cast(pl.String)).collect()
            )
            missing_trees = gh_tags.join(
                trees.select(col("tag").unique()), on="tag", how="anti"
            )
            if missing_trees.is_empty():
                print(f"Already up-to-date {fp!s}")
                result = trees
            else:
                fresh = self._trees_batched(_iter_rows(missing_trees, stop, SemVerTag))
                result = pl.concat((trees, fresh))
        return (
            result.lazy()
            .with_columns(col("tag").cast(semver.tag_enum(gh_tags)))
            .sort("tag", descending=True)
            .collect()
        )

    def refresh_tags(self, npm_tags: pl.DataFrame, /) -> pl.DataFrame:
        limit = self.rate_limit(strict=True)
        npm_tag_only = npm_tags.lazy().select("tag")
        fp = self._paths["tags"]
        if not limit["is_auth"] and limit["remaining"] <= self.req._TAGS_COST:
            return pl.scan_parquet(fp).join(npm_tag_only, on="tag").collect()
        elif not fp.exists():
            print(f"Initializing {fp!s}")
            tags = self.tags(npm_tags=npm_tag_only)
            print(f"Collected {tags.height} new tags")
            return tags
        else:
            print("Checking for new tags")
            prev = pl.scan_parquet(fp)
            latest = self.tags(1, npm_tags=npm_tag_only)
            if latest.equals(prev.pipe(semver.sort).head(1).collect()):
                print(f"Already up-to-date {fp!s}")
                return prev.collect()
            print(f"Refreshing {fp!s}")
            prev_eager = prev.collect()
            tags = (
                pl.concat((self.tags(npm_tags=npm_tag_only), prev_eager))
                .unique("sha")
                .pipe(semver.sort)
            )
            print(f"Collected {tags.height - prev_eager.height} new tags")
            return tags

    def _trees_batched(self, tags: Iterable[str | ParsedTag], /) -> pl.DataFrame:
        rate_limit = self.rate_limit(strict=True)
        if not isinstance(tags, Sequence):
            tags = tuple(tags)
        req = self.req
        n = len(tags)
        cost = req._TREES_COST * n
        if rate_limit["remaining"] < cost:
            raise NotImplementedError(rate_limit, cost)
        print(
            f"Collecting metadata for {n} missing releases.\n"
            f"Using {self.delay(rate_limit):.2f}[ms] between requests ..."
        )
        dfs: list[pl.DataFrame] = []
        for tag in tags:
            time.sleep(self.delay(rate_limit))
            dfs.append(self.trees(tag))
        df = pl.concat(dfs)
        print(f"Finished collection.\nFound {df.height} new rows")
        return df


def _iter_rows(df: pl.DataFrame, stop: int | None, /, tp: type[_TD]) -> Iterator[_TD]:
    """
    Wraps `pl.DataFrame.iter_rows`_ with typing to preserve key completions.

    Parameters
    ----------
    df
        Target dataframe.
    stop
        Passed to `itertools.islice`_.
    tp
        Static type representing a row/record.

        .. note::
            Performs a **very basic** runtime check on the type of ``tp`` (*not* ``df``).

            Primarily used to override ``dict[str, Any]`` when a *narrower* type is known.

    .. _itertools.islice:
        https://docs.python.org/3/library/itertools.html#itertools.islice
    .. _pl.DataFrame.iter_rows:
        https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.iter_rows.html
    """
    if not TYPE_CHECKING:
        assert is_typeddict(tp) or issubclass(tp, Mapping)

    return cast("Iterator[_TD]", islice(df.iter_rows(named=True), stop))