mention_personal_account_in_comments.py

"""
GitHub doesn't currently support updating all comments made on one account
to be made by another account. I made a lot of comments with my work
account, and people keep tagging my work account on issues, but I'd like
them to start referencing my personal account.

So for now, I'll just get every comment I made on my work account and add
a blurb to the beginning about using my personal account.

Run with:

    python3 mention_personal_account_in_comments.py --token=<token> <personal_account>

The token must be a classic token on the work account, and must have
the `public_repo` permissions. Fine-grained personal access tokens do
not currently seem to work.
"""

from __future__ import annotations

import argparse
import dataclasses
import difflib
import enum
import http.client
import itertools
import json
import re
from typing import Any, Iterator, NamedTuple

UPDATE_INDICATOR = "<!-- updated by mention_personal_account_in_comments.py -->"

def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("personal_account")
    parser.add_argument("--fixup-from")
    parser.add_argument("--token")
    parser.add_argument("--limit", type=int)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    personal_account = args.personal_account
    token = args.token
    limit = args.limit
    dry_run = args.dry_run
    fixup_pattern = re.compile(f"@{args.fixup_from}")

    gh = GitHubApi(token=token, user_agent=f"{personal_account} on behalf of work account")

    comments = itertools.chain(
        GitHubComment.load_all(gh),
        GitHubPR.load_all(gh),
    )

    for comment in itertools.islice(comments, limit):
        print("")
        print("*" * 80)
        print(f"***** {comment.url}")
        print("*" * 80)

        if comment.cant_update_reasons:
            print(f"Can't update: {comment.cant_update_reasons}")
            continue

        if UPDATE_INDICATOR not in comment.body:
            update_type = UpdateType.NEEDS_MENTION
        elif fixup_pattern.search(comment.body) is not None:
            update_type = UpdateType.NEEDS_FIXUP
        else:
            update_type = UpdateType.NONE

        if update_type == UpdateType.NONE:
            print("Already updated.")
            continue

        old_body = comment.body
        if update_type == UpdateType.NEEDS_MENTION:
            new_body = "\n".join([
                f":sparkles: _**This is an old work account. Please reference `@{personal_account}` for all future communication**_ :sparkles:",
                UPDATE_INDICATOR,
                "",
                "---",
                "",
                old_body,
            ])
        elif update_type == UpdateType.NEEDS_FIXUP:
            new_body = fixup_pattern.sub(f"@{personal_account}", old_body)
        else:
            raise ValueError(f"Unknown update type: {update_type}")

        if dry_run:
            print("[dry run] Would have made the following changes:")
            print(get_diff(old_body, new_body))
        else:
            comment.update(gh, new_body)
            print("Updated.")

class UpdateType(enum.Enum):
    NEEDS_MENTION = "NEEDS_MENTION"
    NEEDS_FIXUP = "NEEDS_FIXUP"
    NONE = "NONE"

@dataclasses.dataclass
class UpdatableComment:
    url: str
    body: str
    cant_update_reasons: list[str]

    def update(self, gh: GitHubApi, new_body: str) -> None:
        raise NotImplemented

@dataclasses.dataclass
class GitHubComment(UpdatableComment):
    id: int
    repo_owner: str
    repo_name: str

    @classmethod
    def load_all(cls, gh: GitHubApi) -> Iterator[GitHubComment]:
        comments = gh.query_graphql_paginated(
            query="""
            query ($after: String) {
                viewer {
                    issueComments(after: $after, first: 20) {
                        pageInfo {
                            endCursor
                            hasNextPage
                        }
                        nodes {
                            url
                            body
                            databaseId
                            repository {
                                owner {
                                    login
                                }
                                name
                            }
                            viewerCannotUpdateReasons
                        }
                    }
                }
            }
            """,
            connection="data.viewer.issueComments",
        )
        for comment in comments:
            yield cls(
                url=comment["url"],
                body=comment["body"],
                cant_update_reasons=comment["viewerCannotUpdateReasons"],
                id=comment["databaseId"],
                repo_owner=comment["repository"]["owner"]["login"],
                repo_name=comment["repository"]["name"],
            )

    def update(self, gh: GitHubApi, new_body: str) -> None:
        endpoint = f"/repos/{self.repo_owner}/{self.repo_name}/issues/comments/{self.id}"
        print(f">>> Sending request to: {endpoint}")
        gh.query_restapi("PATCH", endpoint, {"body": new_body})

@dataclasses.dataclass
class GitHubPR(UpdatableComment):
    number: int
    repo_owner: str
    repo_name: str

    @classmethod
    def load_all(cls, gh: GitHubApi) -> Iterator[GitHubPR]:
        prs = gh.query_graphql_paginated(
            query="""
            query ($after: String) {
                viewer {
                    pullRequests(after: $after, first: 20) {
                        pageInfo {
                            endCursor
                            hasNextPage
                        }
                        nodes {
                            url
                            body
                            number
                            repository {
                                owner {
                                    login
                                }
                                name
                            }
                            viewerCannotUpdateReasons
                        }
                    }
                }
            }
            """,
            connection="data.viewer.pullRequests",
        )
        for pr in prs:
            yield cls(
                url=pr["url"],
                body=pr["body"],
                cant_update_reasons=pr["viewerCannotUpdateReasons"],
                number=pr["number"],
                repo_owner=pr["repository"]["owner"]["login"],
                repo_name=pr["repository"]["name"],
            )

    def update(self, gh: GitHubApi, new_body: str) -> None:
        endpoint = f"/repos/{self.repo_owner}/{self.repo_name}/pulls/{self.number}"
        print(f">>> Sending request to: {endpoint}")
        gh.query_restapi("PATCH", endpoint, {"body": new_body})

class GitHubApi:
    def __init__(self, *, token: str, user_agent: str) -> None:
        self._token = token
        self._user_agent = user_agent
        self._conn = http.client.HTTPSConnection("api.github.com")

    def query_graphql_paginated(
        self,
        query: str,
        *,
        variables: dict[str, Any] = {},
        connection: str,
    ) -> Iterator[dict[str, Any]]:
        last_cursor = None
        has_next = True
        while has_next:
            response = self.query_graphql(
                query=query,
                variables={
                    "after": last_cursor,
                    **variables,
                },
            )
            for k in connection.split("."):
                response = response[k]

            yield from response["nodes"]

            last_cursor = response["pageInfo"]["endCursor"]
            has_next = response["pageInfo"]["hasNextPage"]

    def query_graphql(self, query: str, variables: dict[str, Any] = {}) -> dict[str, Any]:
        return self._send_request(
            "POST",
            "/graphql",
            body={
                "query": query,
                "variables": variables,
            },
        )

    def query_restapi(self, method: str, endpoint: str, body: dict[str, Any] | None = None) -> dict[str, Any]:
        return self._send_request(
            method,
            endpoint,
            headers={
                "Accept": "application/vnd.github+json",
                "X-GitHub-Api-Version": "2022-11-28",
            },
            body=body,
        )

    def _send_request(
        self,
        method: str,
        endpoint: str,
        *,
        headers: dict[str, str] = {},
        body: dict[str, Any] | str | None = None,
    ) -> dict[str, Any]:
        self._conn.request(
            method,
            endpoint,
            headers={
                "Authorization": f"bearer {self._token}",
                "User-Agent": self._user_agent,
                **headers,
            },
            body=json.dumps(body) if isinstance(body, dict) else body,
        )

        response = self._conn.getresponse()
        resp_body = response.read()

        if response.status != 200:
            raise Exception(f"Error querying GitHub:\n{resp_body}")

        return json.loads(resp_body)

def get_diff(old: str, new: str) -> str:
    return "\n".join(
        difflib.unified_diff(
            old.splitlines(),
            new.splitlines(),
            fromfile="before",
            tofile="after",
            lineterm="",
        )
    )

if __name__ == "__main__":
    main()