Improve the branch selection in interactive go (now same as status)

Author: PawelLipski

git_machete/client/go_interactive.py

@@ -1,5 +1,5 @@
 import sys
-from typing import List, Optional, Tuple
+from typing import Optional, Tuple
 
 try:
     import termios
@@ -10,20 +10,21 @@
     tty = None  # type: ignore[assignment]
 
 from git_machete import utils
-from git_machete.client.base import MacheteClient
+from git_machete.client.status import (StatusData, StatusFlags,
+                                       StatusMacheteClient)
 from git_machete.exceptions import MacheteException, UnexpectedMacheteException
 from git_machete.git_operations import LocalBranchShortName
 from git_machete.utils import bold, index_or_none, warn
 
 
-class GoInteractiveMacheteClient(MacheteClient):
+class GoInteractiveMacheteClient(StatusMacheteClient):
     """Client for interactive branch selection using curses-style interface (implemented without curses, just using ANSI sequences)."""
 
     MAX_VISIBLE_BRANCHES_DEFAULT = 20
     MAX_VISIBLE_BRANCHES_LOWER = 2
     MAX_VISIBLE_BRANCHES_UPPER = 50
 
-    _managed_branches_with_depths: List[Tuple[LocalBranchShortName, int]]
+    _status_data: StatusData
     _current_branch: Optional[LocalBranchShortName]
     _max_visible_branches: int
 
@@ -37,35 +38,10 @@ def _get_max_visible_branches(self) -> int:
         max_visible_branches = terminal_height - 2
         return max(self.MAX_VISIBLE_BRANCHES_LOWER, min(max_visible_branches, self.MAX_VISIBLE_BRANCHES_UPPER))
 
-    def _get_branch_list_with_depths(self) -> List[Tuple[LocalBranchShortName, int]]:
-        """Get a flat list of branches with their depths using DFS traversal."""
-        result: List[Tuple[LocalBranchShortName, int]] = []
-
-        def add_branch_and_children(branch: LocalBranchShortName, depth: int) -> None:
-            result.append((branch, depth))
-            for child_branch in self.down_branches_for(branch) or []:
-                add_branch_and_children(child_branch, depth + 1)
-
-        for root in self._state.roots:
-            add_branch_and_children(root, depth=0)
-
-        return result
-
-    def _render_branch_line(self, branch: LocalBranchShortName, depth: int) -> str:
-        """Render a single branch line with indentation."""
-        indent = "  " * depth
-        marker = " " if self._current_branch is None or branch != self._current_branch else "*"
-
-        line = f"{indent}{marker} {branch}"
-        annotation = self.annotations.get(branch)
-        if annotation and annotation.formatted_full_text:
-            line += f"  {annotation.formatted_full_text}"
-
-        return line
-
     def _draw_screen(self, *, selected_idx: int, scroll_offset: int,
-                     num_lines_drawn: int, is_first_draw: bool) -> int:
-        """Draw the branch selection screen using ANSI escape codes."""
+                     num_lines_drawn: int, is_first_draw: bool) -> Tuple[int, int]:
+        """Draw the branch selection screen using status-style output (format_status_output).
+        Returns (scroll_offset, actual_lines_drawn) so the caller can move the cursor up correctly on next redraw."""
         # Move cursor up to the start of our display area (if we've drawn before)
         if not is_first_draw and num_lines_drawn > 0:
             sys.stdout.write(f'{utils.AE.CSI}{num_lines_drawn}A')
@@ -79,27 +55,30 @@ def _draw_screen(self, *, selected_idx: int, scroll_offset: int,
                        "Enter or Space: checkout, q or Ctrl+C: quit)")
         sys.stdout.write(bold(header_text) + '\n')
 
-        # Adjust scroll offset if needed
-        visible_lines = min(self._max_visible_branches, len(self._managed_branches_with_depths))
-        if selected_idx < scroll_offset:
-            scroll_offset = selected_idx
-        elif selected_idx >= scroll_offset + visible_lines:
-            scroll_offset = selected_idx - visible_lines + 1
-
-        # Draw branches
-        for i in range(visible_lines):
-            branch_idx = scroll_offset + i
-            branch, depth = self._managed_branches_with_depths[branch_idx]
-            line = self._render_branch_line(branch, depth)
-
-            if branch_idx == selected_idx:
-                # Highlight selected line (inverse video)
-                sys.stdout.write(f'{utils.AE.REVERSE_VIDEO}{line}{utils.AE.ENDC}\n')
-            else:
-                sys.stdout.write(f'{line}\n')
+        branches = self._status_data.branches_in_display_order
+        selected_branch = branches[selected_idx] if 0 <= selected_idx < len(branches) else None
+        formatted = self.format_status_output(
+            self._status_data,
+            selected_branch=selected_branch,
+        )
+        lines = formatted.result.splitlines()
+        num_lines = len(lines)
+        visible_lines = min(self._max_visible_branches, num_lines)
+        # line_for_branch maps branch -> 0-based line index in result
+        selected_line_idx = formatted.line_for_branch.get(selected_branch, 0) if selected_branch else 0
+        if selected_line_idx < scroll_offset:
+            scroll_offset = selected_line_idx
+        elif selected_line_idx >= scroll_offset + visible_lines:
+            scroll_offset = selected_line_idx - visible_lines + 1
+
+        lines_drawn = 1  # header
+        end = min(scroll_offset + visible_lines, num_lines)
+        for line_idx in range(scroll_offset, end):
+            sys.stdout.write(lines[line_idx] + '\n')
+            lines_drawn += 1
 
         sys.stdout.flush()
-        return scroll_offset
+        return scroll_offset, lines_drawn
 
     def _get_stdin_fd(self) -> int:  # pragma: no cover; always mocked in tests
         return sys.stdin.fileno()
@@ -136,28 +115,33 @@ def go_interactive(self, *, current_branch: Optional[LocalBranchShortName]) -> O
         """
         Launch interactive branch selection interface.
         Returns the selected branch or None if cancelled.
+        Status data is computed once at start; only rendering (format_status_output) runs on each redraw.
         """
         if termios is None or tty is None:
             raise UnexpectedMacheteException("Interactive mode is not supported on Windows yet")
         if not utils.is_stdout_a_tty():
             raise MacheteException("Interactive `git machete go` requires stdout to be a TTY.")
 
-        # Get flat list of branches with depths from already-parsed state
-        self._managed_branches_with_depths = self._get_branch_list_with_depths()
-
         self._current_branch = current_branch
-
-        # Determine maximum visible branches from terminal height
         self._max_visible_branches = self._get_max_visible_branches()
 
+        # Compute status data once (no list-commits in TUI; config same as status)
+        flags = StatusFlags(
+            maybe_space_before_branch_name=(' ' if self._config.status_extra_space_before_branch_name() else ''),
+            opt_list_commits=False,
+            opt_list_commits_with_hashes=False,
+            opt_squash_merge_detection=self._config.squash_merge_detection(),
+        )
+        self._status_data = self.compute_status_data(flags=flags)
+        branches_ordered = self._status_data.branches_in_display_order
+
         # Find initial selection (current branch or first branch if detached HEAD)
         if current_branch is not None:
-            selected_idx = index_or_none(self.managed_branches, self._current_branch)
+            selected_idx = index_or_none(branches_ordered, self._current_branch)
             if selected_idx is None:
                 selected_idx = 0
                 warn(f"current branch {self._current_branch} is unmanaged\n")
         else:
-            # Detached HEAD - start with first managed branch
             selected_idx = 0
 
         scroll_offset = 0
@@ -168,55 +152,43 @@ def go_interactive(self, *, current_branch: Optional[LocalBranchShortName]) -> O
         sys.stdout.write(utils.AE.HIDE_CURSOR)
         sys.stdout.flush()
 
+        branches = self._status_data.branches_in_display_order
         try:
             while True:
-                # Calculate how many lines we'll draw (header + visible branches)
-                visible_lines = min(self._max_visible_branches, len(self._managed_branches_with_depths))
-                num_lines_drawn = visible_lines + 1  # +1 for header
-
-                scroll_offset = self._draw_screen(
+                scroll_offset, num_lines_drawn = self._draw_screen(
                     selected_idx=selected_idx,
                     scroll_offset=scroll_offset,
                     num_lines_drawn=num_lines_drawn,
                     is_first_draw=is_first_draw
                 )
                 is_first_draw = False
 
-                # Read key
                 key = self._getch()
 
                 if key == utils.AE.KEY_UP:
-                    # Wrap around from first to last
-                    selected_idx = (selected_idx - 1) % len(self._managed_branches_with_depths)
+                    selected_idx = (selected_idx - 1) % len(branches)
                 elif key == utils.AE.KEY_DOWN:
-                    # Wrap around from last to first
-                    selected_idx = (selected_idx + 1) % len(self._managed_branches_with_depths)
+                    selected_idx = (selected_idx + 1) % len(branches)
                 elif key == utils.AE.KEY_SHIFT_UP:
-                    # Jump to first branch
                     selected_idx = 0
                 elif key == utils.AE.KEY_SHIFT_DOWN:
-                    # Jump to last branch
-                    selected_idx = len(self._managed_branches_with_depths) - 1
+                    selected_idx = len(branches) - 1
                 elif key == utils.AE.KEY_LEFT:
-                    # Go to parent
-                    selected_branch, _ = self._managed_branches_with_depths[selected_idx]
+                    selected_branch = branches[selected_idx]
                     parent_branch = self.up_branch_for(selected_branch)
-                    if parent_branch:
-                        selected_idx = self.managed_branches.index(parent_branch)
+                    if parent_branch is not None:
+                        selected_idx = branches.index(parent_branch)
                 elif key == utils.AE.KEY_RIGHT:
-                    # Go to first child
-                    selected_branch, _ = self._managed_branches_with_depths[selected_idx]
+                    selected_branch = branches[selected_idx]
                     child_branches = self.down_branches_for(selected_branch)
                     if child_branches:
-                        selected_idx = self.managed_branches.index(child_branches[0])
+                        selected_idx = branches.index(child_branches[0])
                 elif key in utils.AE.KEYS_ENTER or key == utils.AE.KEY_SPACE:
-                    selected_branch, _ = self._managed_branches_with_depths[selected_idx]
-                    return selected_branch
+                    return branches[selected_idx]
                 elif key in ('q', 'Q'):
                     return None
                 elif key == utils.AE.KEY_CTRL_C:
                     return None
         finally:
-            # Show cursor again and move past our interface
             sys.stdout.write(utils.AE.SHOW_CURSOR)
             sys.stdout.flush()

git_machete/client/status.py

@@ -69,12 +69,26 @@ class StatusData(NamedTuple):
     ongoing_operation: StatusOngoingOperation
 
 
+class StatusFormatOutput(NamedTuple):
+    """Result of formatting status output. Returned by format_status_output."""
+
+    result: str
+    # Maps each branch to the 0-based index of the line in result (when split by newlines) where it appears.
+    line_for_branch: Dict[LocalBranchShortName, int]
+
+
 class StatusMacheteClient(MacheteClient):
     """Client for the status command. Exposes status() and can be used as a mixin for other clients."""
 
     @staticmethod
-    def _format_status_output(data: StatusData) -> str:
-        """Pure function: given StatusData, returns the formatted status tree string."""
+    def format_status_output(
+        data: StatusData,
+        *,
+        selected_branch: Optional[LocalBranchShortName] = None,
+    ) -> StatusFormatOutput:
+        """Pure function: given StatusData, returns StatusFormatOutput (result string and line_for_branch).
+        When selected_branch is set, that branch's name (not annotation/sync status) is wrapped in reverse video.
+        line_for_branch maps each branch to the 0-based line index in result where it appears."""
 
         # These maps need to be defined in a local scope to avoid for mocking the color palette more easily.
         sync_to_parent_status_to_edge_color_map: Dict[SyncToParentStatus, str] = {
@@ -89,9 +103,10 @@ def _format_status_output(data: StatusData) -> str:
             SyncToParentStatus.OUT_OF_SYNC: "x-",
             SyncToParentStatus.MERGED_TO_PARENT: "m-"
         }
-
         out = io.StringIO()
         space = data.flags.maybe_space_before_branch_name
+        line_for_branch: Dict[LocalBranchShortName, int] = {}
+        line_index = 0
 
         next_sibling_of_ancestor_by_branch: Dict[LocalBranchShortName, List[Optional[LocalBranchShortName]]] = {}
 
@@ -127,12 +142,14 @@ def write_line_prefix(
             next_sibling_of_ancestor = next_sibling_of_ancestor_by_branch[branch]
             if b.up_branch is not None:
                 write_line_prefix(branch, next_sibling_of_ancestor, f"{utils.get_vertical_bar()}\n")
+                line_index += 1
                 for commit, fp_suffix in b.commits:
                     write_line_prefix(branch, next_sibling_of_ancestor, utils.get_vertical_bar())
                     out.write(
                         f' {f"{dim(commit.short_hash)}  " if data.flags.opt_list_commits_with_hashes else ""}'
                         f'{dim(commit.subject)}{fp_suffix}\n'
                     )
+                    line_index += 1
                 if utils.ascii_only:
                     junction = sync_to_parent_status_to_junction_ascii_only_map[b.sync_to_parent_status]
                 else:
@@ -145,8 +162,10 @@ def write_line_prefix(
             else:
                 if branch != data.roots[0]:
                     out.write("\n")
+                    line_index += 1
                 out.write("  " + space)
 
+            line_for_branch[branch] = line_index
             op = data.ongoing_operation
             if branch in (op.currently_checked_out_branch, op.currently_rebased_branch, op.currently_bisected_branch):
                 if branch == op.currently_rebased_branch:
@@ -171,9 +190,14 @@ def write_line_prefix(
             if b.annotation is not None and b.annotation.formatted_full_text:
                 anno = '  ' + b.annotation.formatted_full_text
 
-            out.write(current + anno + b.sync_status + b.hook_output + "\n")
+            if selected_branch is not None and branch == selected_branch:
+                current_part = f"{utils.AE.REVERSE_VIDEO}{current}{utils.AE.ENDC}"
+            else:
+                current_part = current
+            out.write(f"{current_part}{anno}{b.sync_status}{b.hook_output}\n")
+            line_index += 1
 
-        return out.getvalue()
+        return StatusFormatOutput(result=out.getvalue(), line_for_branch=line_for_branch)
 
     @staticmethod
     def _status_warning_message(data: StatusData) -> Optional[str]:
@@ -216,7 +240,7 @@ def _status_warning_message(data: StatusData) -> Optional[str]:
             )
         return f"{first_part}.\n\n{second_part}."
 
-    def _compute_status_data(self, *, flags: StatusFlags) -> StatusData:
+    def compute_status_data(self, *, flags: StatusFlags) -> StatusData:
         managed_branches: List[LocalBranchShortName] = list(self._state.managed_branches)
 
         sync_to_parent_status: Dict[LocalBranchShortName, SyncToParentStatus] = {}
@@ -359,9 +383,9 @@ def status(
             opt_list_commits_with_hashes=opt_list_commits_with_hashes,
             opt_squash_merge_detection=opt_squash_merge_detection,
         )
-        data = self._compute_status_data(flags=flags)
-        status_str = self._format_status_output(data)
-        sys.stdout.write(status_str)
+        data = self.compute_status_data(flags=flags)
+        format_out = self.format_status_output(data)
+        sys.stdout.write(format_out.result)
         if warn_when_branch_in_sync_but_fork_point_off:
             warning_msg = self._status_warning_message(data)
             if warning_msg is not None: