Project management automation for MPS

add_issues_to_mps_project.py

#!/usr/bin/env python3

MPS_PROJECT_NUMBER = 26

GH_GET_PROJECT_ID = """
query ($owner: String = "pytorch", $name: String = "pytorch", $number: Int = 26) {
  repository(name: $name, owner: $owner) {
    projectV2(number: $number) {
      id
    }
  }
}
"""

GH_GET_ISSUES_INFO_FRAGMENT = """
fragment IssueInfo on IssueConnection {
  nodes {
    id
    number
    state
    projectsV2(first: 5) {
      nodes {
        number
      }
      totalCount
    }
  }
  pageInfo {
    hasPreviousPage
    startCursor
  }
}
"""

GH_GET_MPS_LABELLED_ISSUES = GH_GET_ISSUES_INFO_FRAGMENT + """
query {
  repository(name: "pytorch", owner: "pytorch") {
    issues(last: 100, labels: ["module: mps"]) {
      ...IssueInfo
    }
  }
}
"""

GH_GET_MPS_LABELLED_ISSUES_CONT = GH_GET_ISSUES_INFO_FRAGMENT + """
query ($cursor: String!) {
  repository(name: "pytorch", owner: "pytorch") {
    issues(last: 100, labels: ["module: mps"], before: $cursor) {
      ...IssueInfo
    }
  }
}
"""

GH_ADD_PR_TO_PROJECT = """
mutation ($project: ID!, $pr: ID!) {
  addProjectV2ItemById(input: {projectId: $project, contentId: $pr}) {
    item {
      id
    }
  }
}
"""


def gh_graphql(query, **kwargs):
    import os
    import json
    from urllib.request import urlopen, Request
    token = os.getenv("GITHUB_TOKEN")
    headers = {"Accept": "application/vnd.github.v3+json", "Authorization": f"token {token}"}

    data = json.dumps({"query": query, "variables": kwargs}).encode()
    with urlopen(Request("https://api.github.com/graphql", data=data, headers=headers)) as conn:
        return json.load(conn)


def get_project_id(project_number=26):
    return gh_graphql(GH_GET_PROJECT_ID, number=project_number)["data"]["repository"]["projectV2"]["id"]


def process_issues(issues, project_id) -> None:
    for issue in issues["nodes"]:
        project_nodes = issue["projectsV2"]["nodes"]
        if len(project_nodes) > 0 and project_nodes[0]["number"] == MPS_PROJECT_NUMBER:
            continue
        print("About to add PR ", issue["number"], " to the project")
        print(gh_graphql(GH_ADD_PR_TO_PROJECT, project=project_id, pr=issue["id"]))


def main() -> None:
    project_id = get_project_id(project_number = MPS_PROJECT_NUMBER)
    data = gh_graphql(GH_GET_MPS_LABELLED_ISSUES)
    issues = data["data"]["repository"]["issues"]
    process_issues(issues, project_id)
    while issues["pageInfo"]["hasPreviousPage"]:
        data = gh_graphql(GH_GET_MPS_LABELLED_ISSUES_CONT, cursor=issues["pageInfo"]["startCursor"])
        issues = data["data"]["repository"]["issues"]
        process_issues(issues, project_id)

if __name__ == "__main__":
    main()