Remove synchronous APIs

.flake8

@@ -17,6 +17,9 @@ max-line-length = 120
 # F403 should turn this on for ghstack/, lookup exclude format TODO
 # E704 see https://github.com/PyCQA/flake8/issues/1925
 ignore = E127, E128, E203, E265, E266, E402, E501, E722, P207, P208, W503, C901, F403, F405, E704
+per-file-ignores =
+  # Script tests are executed by the ghstack test harness and allow top-level await.
+  *.py.test: F401,F704
 exclude =
   .git,
   .hg,

conftest.py

@@ -1,10 +1,11 @@
 # mypy: ignore-errors
 
+import asyncio
 import pathlib
-import runpy
 
 import expecttest
 
+import ghstack.async_script
 import ghstack.test_prelude
 
 import pytest
@@ -33,9 +34,12 @@ def __init__(self, *, direct, **kwargs):
         self.direct = direct
 
     def runtest(self):
-        with ghstack.test_prelude.scoped_test(direct=self.direct):
+        asyncio.run(self.aruntest())
+
+    async def aruntest(self):
+        async with ghstack.test_prelude.scoped_test(direct=self.direct):
             expecttest.EDIT_HISTORY.reload_file(self.fspath)
-            runpy.run_path(self.fspath)
+            await ghstack.async_script.run_path(str(self.fspath))
 
     def repr_failure(self, excinfo):
         excinfo.traceback = excinfo.traceback.cut(path=self.fspath)

src/ghstack/action.py

@@ -8,15 +8,15 @@
 import ghstack.shell
 
 
-def main(
+async def main(
     pull_request: str,
     github: ghstack.github.GitHubEndpoint,
     sh: Optional[ghstack.shell.Shell] = None,
     close: bool = False,
 ) -> None:
 
-    params = ghstack.github_utils.parse_pull_request(pull_request)
-    pr_result = github.graphql(
+    params = await ghstack.github_utils.parse_pull_request(pull_request)
+    pr_result = await github.graphql(
         """
         query ($owner: String!, $name: String!, $number: Int!) {
             repository(name: $name, owner: $owner) {
@@ -32,7 +32,7 @@ def main(
 
     if close:
         logging.info("Closing {owner}/{name}#{number}".format(**params))
-        github.graphql(
+        await github.graphql(
             """
             mutation ($input: ClosePullRequestInput!) {
                 closePullRequest(input: $input) {

src/ghstack/async_script.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+
+import argparse
+import ast
+import asyncio
+import inspect
+import sys
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence
+
+
+async def run_path(
+    path: str,
+    *,
+    argv: Optional[Sequence[str]] = None,
+    globals_: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    script_path = str(Path(path))
+    script_argv = [script_path, *(argv or ())]
+    old_argv = sys.argv
+    sys.argv = script_argv
+    try:
+        source = Path(script_path).read_text()
+        namespace: Dict[str, Any] = {
+            "__file__": script_path,
+            "__name__": "__main__",
+            "__package__": None,
+            "__builtins__": __builtins__,
+        }
+        if globals_ is not None:
+            namespace.update(globals_)
+        code = compile(
+            source,
+            script_path,
+            "exec",
+            flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+        )
+        result = eval(code, namespace)
+        if inspect.isawaitable(result):
+            await result
+        return namespace
+    finally:
+        sys.argv = old_argv
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    parser = argparse.ArgumentParser(
+        description="Run a Python script with top-level await support."
+    )
+    parser.add_argument("script")
+    parser.add_argument("args", nargs=argparse.REMAINDER)
+    ns = parser.parse_args(argv)
+    asyncio.run(run_path(ns.script, argv=ns.args))
+
+
+if __name__ == "__main__":
+    main()

src/ghstack/checkout.py

@@ -8,18 +8,18 @@
 import ghstack.shell
 
 
-def main(
+async def main(
     pull_request: str,
     github: ghstack.github.GitHubEndpoint,
     sh: ghstack.shell.Shell,
     remote_name: str,
     same_base: bool = False,
 ) -> None:
 
-    params = ghstack.github_utils.parse_pull_request(
+    params = await ghstack.github_utils.parse_pull_request(
         pull_request, sh=sh, remote_name=remote_name
     )
-    head_ref = github.get_head_ref(**params)
+    head_ref = await github.get_head_ref(**params)
     orig_ref = re.sub(r"/head$", "/orig", head_ref)
     if orig_ref == head_ref:
         logging.warning(
@@ -31,7 +31,7 @@ def main(
     # If --same-base is specified, check if checkout would change the merge-base
     if same_base:
         # Get the default branch name from the repo
-        repo_info = ghstack.github_utils.get_github_repo_info(
+        repo_info = await ghstack.github_utils.get_github_repo_info(
             github=github,
             sh=sh,
             repo_owner=params["owner"],
@@ -43,24 +43,24 @@ def main(
         default_branch_ref = f"{remote_name}/{default_branch}"
 
         # Get current merge-base with default branch
-        current_base = sh.git("merge-base", default_branch_ref, "HEAD")
+        current_base = await sh.agit("merge-base", default_branch_ref, "HEAD")
     else:
         current_base = None
         default_branch_ref = None
 
-    sh.git("fetch", "--prune", remote_name)
+    await sh.agit("fetch", "--prune", remote_name)
 
     # If --same-base is specified, check what the new merge-base would be
     if same_base:
         assert default_branch_ref is not None
         assert current_base is not None
         target_ref = remote_name + "/" + orig_ref
-        new_base = sh.git("merge-base", default_branch_ref, target_ref)
+        new_base = await sh.agit("merge-base", default_branch_ref, target_ref)
 
         if current_base != new_base:
             raise RuntimeError(
                 f"Checkout would change merge-base from {current_base[:8]} to {new_base[:8]}, "
                 f"aborting due to --same-base flag"
             )
 
-    sh.git("checkout", remote_name + "/" + orig_ref)
+    await sh.agit("checkout", remote_name + "/" + orig_ref)

src/ghstack/cherry_pick.py

@@ -8,7 +8,7 @@
 import ghstack.shell
 
 
-def main(
+async def main(
     pull_request: str,
     github: ghstack.github.GitHubEndpoint,
     sh: ghstack.shell.Shell,
@@ -17,25 +17,25 @@ def main(
     no_fetch: bool = False,
 ) -> None:
 
-    params = ghstack.github_utils.parse_pull_request(
+    params = await ghstack.github_utils.parse_pull_request(
         pull_request, sh=sh, remote_name=remote_name
     )
-    head_ref = github.get_head_ref(**params)
+    head_ref = await github.get_head_ref(**params)
     orig_ref = re.sub(r"/head$", "/orig", head_ref)
     if orig_ref == head_ref:
         logging.warning(
             "The ref {} doesn't look like a ghstack reference".format(head_ref)
         )
 
     if not no_fetch:
-        sh.git("fetch", "--prune", remote_name)
+        await sh.agit("fetch", "--prune", remote_name)
 
     if stack:
         # Cherry-pick the entire stack from merge-base to the commit
         remote_orig_ref = remote_name + "/" + orig_ref
 
         # Find the merge-base with the main branch
-        repo_info = ghstack.github_utils.get_github_repo_info(
+        repo_info = await ghstack.github_utils.get_github_repo_info(
             github=github,
             sh=sh,
             github_url=params["github_url"],
@@ -44,11 +44,11 @@ def main(
         main_branch = f"{remote_name}/{repo_info['default_branch']}"
 
         # Get merge-base between the commit and main branch
-        merge_base = sh.git("merge-base", main_branch, remote_orig_ref).strip()
+        merge_base = (await sh.agit("merge-base", main_branch, remote_orig_ref)).strip()
 
         # Get all commits from merge-base to the target commit
         commit_list = (
-            sh.git("rev-list", "--reverse", f"{merge_base}..{remote_orig_ref}")
+            (await sh.agit("rev-list", "--reverse", f"{merge_base}..{remote_orig_ref}"))
             .strip()
             .split("\n")
         )
@@ -58,10 +58,10 @@ def main(
 
         logging.info(f"Cherry-picking {len(commit_list)} commits from stack")
         for commit in commit_list:
-            sh.git("cherry-pick", commit)
+            await sh.agit("cherry-pick", commit)
             logging.info(f"Cherry-picked {commit}")
     else:
         # Cherry-pick just the single commit
         remote_orig_ref = remote_name + "/" + orig_ref
-        sh.git("cherry-pick", remote_orig_ref)
+        await sh.agit("cherry-pick", remote_orig_ref)
         logging.info(f"Cherry-picked {orig_ref}")