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

Author: PawelLipski

git_machete/client/go_interactive.py

@@ -9,21 +9,35 @@
     termios = None  # type: ignore[assignment]
     tty = None  # type: ignore[assignment]
 
-from git_machete import utils
-from git_machete.client.base import MacheteClient
+from git_machete import git_config_keys, utils
+from git_machete.client.base import SquashMergeDetection
+from git_machete.client.status import (StatusData, StatusFlags,
+                                       StatusMacheteClient)
 from git_machete.exceptions import UnexpectedMacheteException
 from git_machete.git_operations import LocalBranchShortName
 from git_machete.utils import AnsiEscapeCodes, bold, index_or_none, warn
 
 
-class GoInteractiveMacheteClient(MacheteClient):
+def _branch_line_indices(data: StatusData, branches: List[LocalBranchShortName]) -> List[int]:
+    """Line index in status output for each branch (status has an extra newline before non-first roots)."""
+    roots = set(data.roots)
+    first_root = next((b for b in branches if b in roots), None)
+    result: List[int] = []
+    for i, branch in enumerate(branches):
+        blanks_before = sum(1 for j in range(i) if branches[j] in roots and branches[j] != first_root)
+        result.append(i + blanks_before)
+    return result
+
+
+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
+    _branch_line_indices: List[int]  # line index in status output for each branch
     _current_branch: Optional[LocalBranchShortName]
     _max_visible_branches: int
 
@@ -37,35 +51,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'{AnsiEscapeCodes.CSI}{num_lines_drawn}A')
@@ -79,27 +68,31 @@ 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
+        branches = self._status_data.branches_in_display_order
+        line_indices = self._branch_line_indices
+        selected_branch = branches[selected_idx] if 0 <= selected_idx < len(branches) else None
+        status_str = self.format_status_output(
+            self._status_data,
+            selected_branch=selected_branch,
+        )
+        lines = status_str.splitlines()
+        num_lines = len(lines)
+        visible_lines = min(self._max_visible_branches, num_lines)
+        selected_line_idx = line_indices[selected_idx] if 0 <= selected_idx < len(line_indices) 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
         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'{AnsiEscapeCodes.REVERSE_VIDEO}{line}{AnsiEscapeCodes.ENDC}\n')
-            else:
-                sys.stdout.write(f'{line}\n')
+            line_idx = scroll_offset + i
+            if line_idx < num_lines:
+                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,26 +129,36 @@ 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")
 
-        # 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; squash merge detection defaults to simple)
+        maybe_space = (
+            ' ' if self._git.get_boolean_config_attr(
+                git_config_keys.STATUS_EXTRA_SPACE_BEFORE_BRANCH_NAME, default_value=False) else ''
+        )
+        flags = StatusFlags(
+            maybe_space_before_branch_name=maybe_space,
+            opt_list_commits=False,
+            opt_list_commits_with_hashes=False,
+            opt_squash_merge_detection=SquashMergeDetection.SIMPLE,
+        )
+        self._status_data = self.compute_status_data(flags=flags)
+        branches_ordered = self._status_data.branches_in_display_order
+        self._branch_line_indices = _branch_line_indices(self._status_data, branches_ordered)
+
         # 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
@@ -166,55 +169,43 @@ def go_interactive(self, *, current_branch: Optional[LocalBranchShortName]) -> O
         sys.stdout.write(AnsiEscapeCodes.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 == AnsiEscapeCodes.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 == AnsiEscapeCodes.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 == AnsiEscapeCodes.KEY_SHIFT_UP:
-                    # Jump to first branch
                     selected_idx = 0
                 elif key == AnsiEscapeCodes.KEY_SHIFT_DOWN:
-                    # Jump to last branch
-                    selected_idx = len(self._managed_branches_with_depths) - 1
+                    selected_idx = len(branches) - 1
                 elif key == AnsiEscapeCodes.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 == AnsiEscapeCodes.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 AnsiEscapeCodes.KEYS_ENTER or key == AnsiEscapeCodes.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 == AnsiEscapeCodes.KEY_CTRL_C:
                     return None
         finally:
-            # Show cursor again and move past our interface
             sys.stdout.write(AnsiEscapeCodes.SHOW_CURSOR)
             sys.stdout.flush()

git_machete/client/status.py

@@ -87,8 +87,13 @@ 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,
+    ) -> str:
+        """Pure function: given StatusData, returns the formatted status tree string.
+        When selected_branch is set, that branch's name (not annotation/sync status) is wrapped in reverse video."""
         out = io.StringIO()
         space = data.flags.maybe_space_before_branch_name
 
@@ -169,7 +174,11 @@ 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"{AnsiEscapeCodes.REVERSE_VIDEO}{current}{AnsiEscapeCodes.ENDC}"
+            else:
+                current_part = current
+            out.write(f"{current_part}{anno}{b.sync_status}{b.hook_output}\n")
 
         return out.getvalue()
 
@@ -214,7 +223,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] = {}
@@ -360,8 +369,8 @@ 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)
+        data = self.compute_status_data(flags=flags)
+        status_str = self.format_status_output(data)
         sys.stdout.write(status_str)
         if warn_when_branch_in_sync_but_fork_point_off:
             warning_msg = self._status_warning_message(data)