Infer direct mode from existing pull requests

AGENTS.md

@@ -2,7 +2,7 @@
 
 ## Lint
 
-Run `python -m black src/ghstack/` and `python -m flake8 src/ghstack/` to fix and check lint before committing. If lint fixes are needed after a commit, amend the commit rather than creating a separate lint fix commit.
+Run `uv run --frozen lintrunner --all-files` to reproduce CI lint locally before committing. Invoke lintrunner through `uv run` unless the project virtual environment is active; otherwise its Python-based linters may use a system Python without the development dependencies. To apply formatter fixes, run `uv run --frozen lintrunner format --all-files`, then rerun the full lint command. If lint fixes are needed after a commit, amend the commit rather than creating a separate lint fix commit.
 
 ## Tests
 

conftest.py

@@ -17,15 +17,29 @@
 def pytest_collect_file(file_path: pathlib.Path, parent):
     # NB: script name must not end with py, due to doctest picking it
     # up in that case
-    if file_path.suffixes == [".py", ".test"]:
+    if file_path.name.endswith(".py.test"):
         return Script.from_parent(parent, path=file_path)
 
 
 class Script(pytest.File):
     def collect(self):
-        yield ScriptItem.from_parent(self, name="default", direct=False)
-        if self.path.parent.name in ["submit", "unlink", "log", "sync"]:
-            yield ScriptItem.from_parent(self, name="direct", direct=True)
+        if self.path.name.endswith(".unparametrized.py.test"):
+            direct_values = [None]
+        else:
+            direct_values = [False]
+            if self.path.parent.name in ["submit", "unlink", "log", "sync"]:
+                direct_values.append(True)
+
+        for direct in direct_values:
+            yield ScriptItem.from_parent(
+                self,
+                name=(
+                    "unparametrized"
+                    if direct is None
+                    else "direct" if direct else "default"
+                ),
+                direct=direct,
+            )
 
 
 class ScriptItem(pytest.Item):

src/ghstack/submit.py

@@ -516,21 +516,33 @@ async def initialize(self) -> None:
 
         object.__setattr__(self, "base", default_branch)
 
-        # Check if direct should be used, if the user didn't explicitly
-        # specify an option
+    # ~~~~~~~~~~~~~~~~~~~~~~~~
+    # The main algorithm
+
+    async def _initialize_direct(self, pr_info_cache: Dict[GitHubNumber, Any]) -> None:
         direct = self.direct_opt
         if direct is None:
             direct_r = await self.sh.agit(
                 "cat-file", "-e", "HEAD:.github/ghstack_direct", exitcode=True
             )
             assert isinstance(direct_r, bool)
             direct = direct_r
-
+        if self.direct_opt is None and not direct:
+            styles = {
+                re.fullmatch(r"gh/[^/]+/[0-9]+/base", base_ref) is None
+                for pr_info in pr_info_cache.values()
+                if (base_ref := self._pr_ref_name(pr_info, "base")) is not None
+            }
+            if len(styles) > 1:
+                raise RuntimeError(
+                    "Cannot infer ghstack submission style: the stack contains "
+                    "both direct and non-direct pull requests. Pass --direct or "
+                    "--no-direct explicitly."
+                )
+            if styles:
+                direct = styles.pop()
         object.__setattr__(self, "direct", direct)
 
-    # ~~~~~~~~~~~~~~~~~~~~~~~~
-    # The main algorithm
-
     async def run(self) -> List[DiffMeta]:
         timer = _Timer() if _TIMING_ENABLED else None
 
@@ -601,6 +613,7 @@ async def run(self) -> List[DiffMeta]:
         )
 
         pr_info_cache = await self._prefetch_pr_info(commits_to_rebase)
+        await self._initialize_direct(pr_info_cache)
         if not self.no_fetch:
             await self._fetch_foreign_pr_refs(pr_info_cache.values())
 

src/ghstack/test_prelude.py

@@ -122,9 +122,9 @@ class Context:
     github: ghstack.github.GitHubEndpoint
     upstream_sh: ghstack.shell.Shell
     sh: ghstack.shell.Shell
-    direct: bool
+    direct: Optional[bool]
 
-    def __init__(self, direct: bool) -> None:
+    def __init__(self, direct: Optional[bool]) -> None:
         # Set up a "parent" repository with an empty initial commit that we'll operate on
         upstream_dir = tempfile.mkdtemp()
         self.upstream_sh = ghstack.shell.Shell(cwd=upstream_dir, testing=True)
@@ -154,7 +154,7 @@ def cleanup(self) -> None:
                 onerror=handle_remove_read_only,
             )
 
-    async def check_global_github_invariants(self, direct: bool) -> None:
+    async def check_global_github_invariants(self, direct: Optional[bool]) -> None:
         r = await self.github.graphql(
             """
           query {
@@ -177,7 +177,7 @@ async def check_global_github_invariants(self, direct: bool) -> None:
                 continue
             # In direct mode, only head refs may not be reused;
             # base refs can be reused in octopus situations
-            if not direct:
+            if direct is False:
                 assert pr["baseRefName"] not in seen_refs
                 seen_refs.add(pr["baseRefName"])
             assert pr["headRefName"] not in seen_refs
@@ -200,7 +200,7 @@ async def init_test() -> Context:
 
 
 @contextlib.asynccontextmanager
-async def scoped_test(direct: bool) -> AsyncIterator[None]:
+async def scoped_test(direct: Optional[bool]) -> AsyncIterator[None]:
     global CTX
     assert CTX is None
     try:
@@ -224,8 +224,10 @@ async def gh_submit(
     reviewer: Optional[str] = None,
     label: Optional[str] = None,
     automsg: Optional[str] = None,
+    **submit_kwargs: Any,
 ) -> List[ghstack.submit.DiffMeta]:
     self = CTX
+    direct_opt = submit_kwargs.pop("direct_opt", self.direct)
     r = await ghstack.submit.main(
         msg=msg,
         username="ezyang",
@@ -236,7 +238,7 @@ async def gh_submit(
         repo_owner_opt="pytorch",
         repo_name_opt="pytorch",
         short=short,
-        direct_opt=self.direct,
+        direct_opt=direct_opt,
         no_skip=no_skip,
         github_url="github.com",
         remote_name="origin",
@@ -247,6 +249,7 @@ async def gh_submit(
         reviewer=reviewer,
         label=label,
         automsg=automsg,
+        **submit_kwargs,
     )
     await self.check_global_github_invariants(self.direct)
     return r
@@ -451,7 +454,7 @@ async def assert_github_state(expect: str, *, skip: int = 0) -> None:
     assert_expected_inline(await dump_github(), expect, skip=skip + 1)
 
 
-def is_direct() -> bool:
+def is_direct() -> Optional[bool]:
     return CTX.direct
 
 

test/submit/infer_direct.unparametrized.py.test

@@ -0,0 +1,11 @@
+from ghstack.test_prelude import *
+
+await init_test()
+
+await commit("A")
+await gh_submit("Initial", direct_opt=True)
+
+await amend("A2")
+await gh_submit("Update")
+
+ok()