Add jlab script

hm/common.nix

@@ -148,6 +148,7 @@
       rename
       which
       file
+      cached-nix-shell # For scripts
 
       # Pipe utils
       gnugrep

hm/git/default.nix

@@ -21,7 +21,7 @@ in
                   then
                     ${lib.getExe config.programs.jujutsu.package} git fetch
                     ${lib.getExe config.programs.jujutsu.package} rebase -d main@origin
-                    ${lib.getExe config.programs.jujutsu.package} branch set main -r @-
+                    ${lib.getExe config.programs.jujutsu.package} bookmark set main -r @-
                     ${lib.getExe config.programs.jujutsu.package} git push
                   else
                     ${pkgs.git}/bin/git --no-optional-locks diff --quiet || echo "Repository is dirty!"

hm/scripts/jlab

@@ -0,0 +1,265 @@
+#!/usr/bin/env cached-nix-shell
+#! nix-shell -i python3
+#! nix-shell -p python3 python3Packages.pydantic
+# vim: filetype=python
+
+"""
+glab wrapper for jujutsu,
+with some convinience features.
+"""
+
+import re
+import subprocess
+import sys
+import typing
+
+import pydantic
+import typing_extensions
+
+
+class GitLabMR(pydantic.BaseModel):
+    """
+    Represents a GitLab Merge Request.
+    """
+
+    title: str
+    source_branch: str
+    target_branch: str
+    project_id: int
+    source_project_id: int
+    target_project_id: int
+
+    @pydantic.model_validator(mode="after")
+    def same_project(self) -> typing_extensions.Self:
+        if not (self.project_id == self.source_project_id == self.target_project_id):
+            raise NotImplementedError("Different project ids")
+        return self
+
+
+def glab_get_mr(branch: str) -> GitLabMR:
+    """
+    Get details about a GitLab MR.
+    """
+    sp = subprocess.run(
+        ["glab", "mr", "view", branch, "--output", "json"], stdout=subprocess.PIPE
+    )
+    sp.check_returncode()
+    return GitLabMR.model_validate_json(sp.stdout)
+
+
+class JujutsuType:
+    """
+    Utility to work with Template types.
+    https://martinvonz.github.io/jj/latest/templates/
+    """
+
+    FIELD_SEPARATOR: typing.ClassVar[str] = "\0"
+    ESCAPED_SEPARATOR: typing.ClassVar[str] = r"\0"
+
+    @staticmethod
+    def template(base: str, type_: typing.Type) -> str:
+        """
+        Generates a --template string that is machine-parseable for a given type.
+        """
+        if typing.get_origin(type_) == list:
+            # If we have a list, prepend it with the number of items
+            # so we know how many fields we should consume.
+            (subtype,) = typing.get_args(type_)
+            subtype = typing.cast(typing.Type, subtype)
+            return (
+                f'{base}.len()++"{JujutsuType.ESCAPED_SEPARATOR}"'
+                f'++{base}.map(|l| {JujutsuType.template("l", subtype)})'
+            )
+        elif issubclass(type_, JujutsuObject):
+            return type_.template(base)
+        else:
+            return f'{base}++"{JujutsuType.ESCAPED_SEPARATOR}"'
+
+    @staticmethod
+    def parse(stack: list[str], type_: typing.Type) -> typing.Any:
+        """
+        Unserialize the result of a template to a given type.
+        Needs to be provided the template as a list splitted by the field separator.
+        It will consume the fields it needs.
+        """
+        if typing.get_origin(type_) == list:
+            (subtype,) = typing.get_args(type_)
+            subtype = typing.cast(typing.Type, subtype)
+            len = int(stack.pop(0))
+            return [JujutsuType.parse(stack, subtype) for i in range(len)]
+        elif issubclass(type_, JujutsuObject):
+            return type_.parse(stack)
+        else:
+            return stack.pop(0)
+
+
+class JujutsuObject(pydantic.BaseModel):
+    @classmethod
+    def template(cls, base: str) -> str:
+        temp = []
+        for k, v in cls.model_fields.items():
+            key = f"{base}.{k}()"
+            temp.append(JujutsuType.template(key, v.annotation))
+        return "++".join(temp)
+
+    @classmethod
+    def parse(cls, stack: list[str]) -> typing_extensions.Self:
+        ddict = dict()
+        for k, v in cls.model_fields.items():
+            ddict[k] = JujutsuType.parse(stack, v.annotation)
+        return cls(**ddict)
+
+
+class JujutsuShortestIdPrefix(JujutsuObject):
+    prefix: str
+    rest: str
+
+    @property
+    def full(self) -> str:
+        return self.prefix + self.rest
+
+
+class JujutsuChangeId(JujutsuObject):
+    shortest: JujutsuShortestIdPrefix
+
+    @property
+    def full(self) -> str:
+        return self.shortest.full
+
+
+class JujutsuRefName(JujutsuObject):
+    name: str
+
+
+class JujutsuCommit(JujutsuObject):
+    change_id: JujutsuChangeId
+    bookmarks: list[JujutsuRefName]
+
+
+class Jujutsu:
+    """
+    Represents a Jujutsu repo.
+    Since there's no need for multi-repo, this is just the one in the current directory.
+    """
+
+    def run(self, *args: str, **kwargs: typing.Any) -> subprocess.CompletedProcess:
+        cmd = ["jj"]
+        cmd.extend(args)
+        sp = subprocess.run(cmd, stdout=subprocess.PIPE)
+        sp.check_returncode()
+        return sp
+
+    def log(self, revset: str = "@") -> list[JujutsuCommit]:
+        cmd = [
+            "log",
+            "-r",
+            revset,
+            "--no-graph",
+            "-T",
+            JujutsuCommit.template("self"),
+        ]
+        sp = self.run(*cmd, stdout=subprocess.PIPE)
+        stack = sp.stdout.decode().split(JujutsuType.FIELD_SEPARATOR)
+        assert stack[-1] == "", "No trailing NUL byte"
+        stack.pop()
+        commits = []
+        while len(stack):
+            commits.append(JujutsuCommit.parse(stack))
+        return commits
+
+
+jj = Jujutsu()
+
+
+def current_bookmark() -> JujutsuRefName | None:
+    """
+    Replacement of git's current branch concept working with jj.
+    Needed for commodity features, such as not requiring to type the MR mumber / branch
+    for `glab mr`, or automatically advance the bookmark to the head before pushing.
+    """
+    bookmarks = []
+    for commit in jj.log("reachable(@, trunk()..)"):
+        bookmarks.extend(commit.bookmarks)
+
+    if len(bookmarks) > 1:
+        raise NotImplementedError("Multiple bookmarks on trunk branch")  # TODO
+        # If there's a split in the tree: TBD
+        # If there's no bookmark ahead: the bookmark behind
+        # If there's a bookmark ahead: that one
+        # (needs adjusting of push so it doesn't advance anything then)
+
+    if bookmarks:
+        return bookmarks[0]
+    else:
+        return None
+
+
+def to_glab() -> None:
+    """
+    Pass the remaining arguments to glab.
+    """
+    sp = subprocess.run(["glab"] + sys.argv[1:])
+    sys.exit(sp.returncode)
+
+
+if len(sys.argv) <= 1:
+    to_glab()
+elif sys.argv[1] in ("merge", "mr"):
+    if len(sys.argv) <= 2:
+        to_glab()
+    elif sys.argv[2] == "checkout":
+        # Bypass the original checkout command so it doesn't run git commands.
+        # If there's no commit on the branch, add one with the MR title
+        # so jj has a current bookmark.
+        mr = glab_get_mr(sys.argv[3])
+        jj.run("git", "fetch")
+        if len(JujutsuCommit.log(f"{mr.source_branch} | {mr.target_branch}")) == 1:
+            title = re.sub(r"^(WIP|Draft): ", "", mr.title)
+            jj.run("new", mr.source_branch)
+            jj.run("describe", "-m", title)
+            jj.run("bookmark", "move", mr.source_branch)
+        else:
+            jj.run("bookmark", "edit", mr.source_branch)
+    elif sys.argv[2] in (
+        # If no MR number/branch is given, insert the current bookmark,
+        # as the current branch concept doesn't exist in jj
+        "approve",
+        "approvers",
+        "checkout",
+        "close",
+        "delete",
+        "diff",
+        "issues",
+        "merge",
+        "note",
+        "rebase",
+        "revoke",
+        "subscribe",
+        "todo",
+        "unsubscribe",
+        "update",
+        "view",
+    ):
+        if len(sys.argv) <= 3 or sys.argv[3].startswith("-"):
+            bookmark = current_bookmark()
+            if bookmark:
+                sys.argv.insert(3, bookmark.name)
+        to_glab()
+    else:
+        to_glab()
+elif sys.argv[1] == "push":
+    # Advance the current branch to the head and push
+    bookmark = current_bookmark()
+    if not bookmark:
+        raise RuntimeError("Couldn't find a current branch")
+    heads = jj.log("heads(@::)")
+    if len(heads) != 1:
+        raise RuntimeError("Multiple heads")  # Or none if something goes horribly wrong
+    head = heads[0]
+    jj.run("bookmark", "set", bookmark.name, "-r", head.change_id.full)
+    jj.run("git", "push", "--bookmark", bookmark.name)
+    # TODO Sign https://github.com/martinvonz/jj/issues/4712
+else:
+    to_glab()
+
+# TODO Autocomplete