No bug: ensure ./mach taskgraph --diff returns you to your starting branch when using git. r=ahal

Differential Revision: https://phabricator.services.mozilla.com/D111848

python/mozversioncontrol/mozversioncontrol/__init__.py

@@ -161,6 +161,10 @@ class Repository(object):
         Return None if the hg hash of the base ref could not be calculated.
         """
 
+    @abc.abstractproperty
+    def branch(self):
+        """Current branch or bookmark the checkout has active."""
+
     @abc.abstractmethod
     def get_commit_time(self):
         """Return the Unix time of the HEAD revision."""
@@ -329,6 +333,16 @@ class HgRepository(Repository):
     def base_ref_as_hg(self):
         return self.base_ref
 
+    @property
+    def branch(self):
+        bookmarks_fn = os.path.join(self.path, ".hg", "bookmarks.current")
+        if os.path.exists(bookmarks_fn):
+            with open(bookmarks_fn) as f:
+                bookmark = f.read()
+                return bookmark or None
+
+        return None
+
     def __enter__(self):
         if self._client.server is None:
             # The cwd if the spawned process should be the repo root to ensure
@@ -541,6 +555,10 @@ class GitRepository(Repository):
         except subprocess.CalledProcessError:
             return
 
+    @property
+    def branch(self):
+        return self._run("branch", "--show-current").strip() or None
+
     @property
     def has_git_cinnabar(self):
         try:

taskcluster/mach_commands.py

@@ -539,7 +539,12 @@ class MachCommands(MachCommandBase):
                     print("abort: can't diff taskgraph with dirty working directory")
                     return 1
 
-                cur_ref = vcs.head_ref[:12]
+                # We want to return the working directory to the current state
+                # as best we can after we're done. In all known cases, using
+                # branch or bookmark (which are both available on the VCS object)
+                # as `branch` is preferable to a specific revision.
+                cur_ref = vcs.branch or vcs.head_ref[:12]
+
                 if options["diff"] == "default":
                     base_ref = vcs.base_ref
                 else: