| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| from git import Repo, InvalidGitRepositoryError |
|
|
| BRANCH_COLORS = { |
| "main": "#4C8BF5", |
| "master": "#4C8BF5", |
| } |
| FALLBACK_COLORS = ["#2ECC71", "#E84B3A", "#F39C12", "#9B59B6", "#1ABC9C", "#E67E22"] |
| GHOST_COLOR = "#FFD700" |
| HEAD_RING_COLOR = "#FFFFFF" |
| BG_COLOR = "#0D1117" |
| TEXT_COLOR = "#C9D1D9" |
| X_SPACING = 1.5 |
| Y_SPACING = -0.7 |
|
|
|
|
| def get_branch_color(branch_name: str, index: int) -> str: |
| return BRANCH_COLORS.get(branch_name, FALLBACK_COLORS[index % len(FALLBACK_COLORS)]) |
|
|
|
|
| def load_repo(repo_path: str) -> Repo: |
| return Repo(repo_path) |
|
|
|
|
| def get_commit_graph(repo: Repo) -> list[dict]: |
| seen = set() |
| commits = [] |
| for commit in repo.iter_commits("--all"): |
| if commit.hexsha in seen: |
| continue |
| seen.add(commit.hexsha) |
| commits.append({ |
| "sha": commit.hexsha, |
| "short_sha": commit.hexsha[:7], |
| "message": commit.message.strip().splitlines()[0], |
| "parents": [p.hexsha for p in commit.parents], |
| }) |
| return commits |
|
|
|
|
| def tag_commits_with_branches(repo: Repo, commits: list[dict]) -> list[dict]: |
| sha_to_local: dict[str, list[str]] = {} |
| for branch in repo.branches: |
| tip = branch.commit.hexsha |
| sha_to_local.setdefault(tip, []).append(branch.name) |
|
|
| sha_to_remote: dict[str, list[str]] = {} |
| for remote in repo.remotes: |
| for ref in remote.refs: |
| if ref.name.endswith("/HEAD"): |
| continue |
| tip = ref.commit.hexsha |
| sha_to_remote.setdefault(tip, []).append(ref.name) |
|
|
| head_sha = repo.head.commit.hexsha |
| for commit in commits: |
| commit["local_labels"] = sha_to_local.get(commit["sha"], []) |
| commit["remote_labels"] = sha_to_remote.get(commit["sha"], []) |
| commit["is_head"] = commit["sha"] == head_sha |
| return commits |
|
|
|
|
| def assign_lanes(commits: list[dict]) -> dict[str, int]: |
| lane_by_sha: dict[str, int] = {} |
| active_lanes: list[str | None] = [] |
|
|
| for commit in commits: |
| sha = commit["sha"] |
| parents = commit["parents"] |
|
|
| inherited_lane = None |
| for parent_sha in parents: |
| if parent_sha in lane_by_sha: |
| candidate = lane_by_sha[parent_sha] |
| if active_lanes[candidate] == parent_sha: |
| inherited_lane = candidate |
| break |
|
|
| if inherited_lane is not None: |
| lane = inherited_lane |
| else: |
| try: |
| lane = active_lanes.index(None) |
| except ValueError: |
| lane = len(active_lanes) |
| active_lanes.append(None) |
|
|
| lane_by_sha[sha] = lane |
| if lane < len(active_lanes): |
| active_lanes[lane] = sha |
| else: |
| active_lanes.append(sha) |
|
|
| return lane_by_sha |
|
|
|
|
| def inject_ghost_node(commits: list[dict], head_sha: str, current_branch: str, message: str) -> list[dict]: |
| ghost = { |
| "sha": "GHOST_NEXT", |
| "short_sha": "next", |
| "message": message.strip(), |
| "parents": [head_sha], |
| "local_labels": [f"← {current_branch} (next)"], |
| "remote_labels": [], |
| "is_head": False, |
| } |
| return [ghost] + commits |
|
|
|
|
| def build_figure(repo: Repo, commits: list[dict], current_branch: str, next_commit_message: str) -> plt.Figure: |
| head_sha = repo.head.commit.hexsha |
| ghost_sha = "GHOST_NEXT" |
|
|
| if next_commit_message.strip(): |
| commits = inject_ghost_node(commits, head_sha, current_branch, next_commit_message) |
|
|
| lane_by_sha = assign_lanes(commits) |
| sha_to_index = {c["sha"]: i for i, c in enumerate(commits)} |
|
|
| num_lanes = max(lane_by_sha.values(), default=0) + 1 |
| fig_width = max(12, 4 + num_lanes * 2) |
| fig_height = max(6, len(commits) * 0.6 + 2) |
|
|
| fig, ax = plt.subplots(figsize=(fig_width, fig_height)) |
| fig.patch.set_facecolor(BG_COLOR) |
| ax.set_facecolor(BG_COLOR) |
|
|
| def node_position(sha: str) -> tuple[float, float]: |
| idx = sha_to_index[sha] |
| lane = lane_by_sha[sha] |
| return lane * X_SPACING, idx * Y_SPACING |
|
|
| for commit in commits: |
| x_child, y_child = node_position(commit["sha"]) |
| lane = lane_by_sha[commit["sha"]] |
| edge_color = get_branch_color("", lane) |
|
|
| for parent_sha in commit["parents"]: |
| if parent_sha not in sha_to_index: |
| continue |
| x_parent, y_parent = node_position(parent_sha) |
| is_merge = x_child != x_parent |
| ax.annotate( |
| "", |
| xy=(x_parent, y_parent), |
| xytext=(x_child, y_child), |
| arrowprops=dict( |
| arrowstyle="-|>", |
| color=edge_color, |
| lw=1.8, |
| connectionstyle=f"arc3,rad={'0.3' if is_merge else '0.0'}", |
| ), |
| zorder=2, |
| ) |
|
|
| for commit in commits: |
| x, y = node_position(commit["sha"]) |
| lane = lane_by_sha[commit["sha"]] |
| node_color = get_branch_color("", lane) |
| is_ghost = commit["sha"] == ghost_sha |
|
|
| if is_ghost: |
| ax.plot(x, y, "o", markersize=13, color=GHOST_COLOR, alpha=0.35, zorder=4) |
| ax.plot(x, y, "o", markersize=13, color=GHOST_COLOR, fillstyle="none", |
| markeredgewidth=2, linestyle="--", zorder=5) |
| elif commit["is_head"]: |
| ax.plot(x, y, "o", markersize=15, color=node_color, zorder=4) |
| ax.plot(x, y, "o", markersize=8, color=HEAD_RING_COLOR, zorder=5) |
| else: |
| ax.plot(x, y, "o", markersize=10, color=node_color, zorder=4) |
|
|
| label_text = f"{commit['short_sha']} {commit['message'][:50]}" |
| ax.text(x + 0.35, y, label_text, va="center", ha="left", |
| fontsize=8, color=TEXT_COLOR, fontfamily="monospace", zorder=6) |
|
|
| badge_offset_x = x + 0.35 |
| badge_offset_y = y + 0.22 |
| for label in commit["local_labels"]: |
| badge_bg = "#9E6A03" if label == current_branch else "#238636" |
| ax.text( |
| badge_offset_x, badge_offset_y, |
| f" {label} ", |
| va="center", ha="left", |
| fontsize=7, color="white", fontfamily="monospace", |
| bbox=dict(boxstyle="round,pad=0.25", facecolor=badge_bg, edgecolor="none"), |
| zorder=7, |
| ) |
| badge_offset_x += len(label) * 0.075 + 0.6 |
| for label in commit["remote_labels"]: |
| ax.text( |
| badge_offset_x, badge_offset_y, |
| f" {label} ", |
| va="center", ha="left", |
| fontsize=7, color="white", fontfamily="monospace", |
| bbox=dict(boxstyle="round,pad=0.25", facecolor="#1F6FEB", edgecolor="none"), |
| zorder=7, |
| ) |
| badge_offset_x += len(label) * 0.075 + 0.6 |
|
|
| ax.set_xlim(-0.8, num_lanes * X_SPACING + 6) |
| ax.set_ylim(len(commits) * Y_SPACING - 0.8, 0.8) |
| ax.set_title("Git Branch Graph", color=TEXT_COLOR, fontsize=12, pad=12) |
| ax.axis("off") |
| plt.tight_layout() |
|
|
| return fig |
|
|
|
|
| def build_info(repo: Repo, current_branch: str, next_commit_message: str) -> str: |
| head_sha = repo.head.commit.hexsha |
| staged_files = [item.a_path for item in repo.index.diff("HEAD")] |
| untracked = repo.untracked_files |
| remote_branches = [ |
| ref.name for remote in repo.remotes for ref in remote.refs if not ref.name.endswith("/HEAD") |
| ] |
|
|
| info_lines = [ |
| f"**Current branch:** `{current_branch}`", |
| f"**HEAD:** `{head_sha[:7]}`", |
| ] |
| if remote_branches: |
| info_lines.append("**Remote branches:** " + ", ".join(f"`{r}`" for r in remote_branches)) |
| if staged_files: |
| info_lines.append("**Staged (ready to commit):** " + ", ".join(f"`{f}`" for f in staged_files)) |
| if untracked: |
| info_lines.append("**Untracked:** " + ", ".join(f"`{f}`" for f in untracked[:5])) |
| if next_commit_message.strip(): |
| info_lines.append( |
| f"**Next commit preview:** `{next_commit_message.strip()}` ← will land on `{current_branch}`" |
| ) |
|
|
| return "\n\n".join(info_lines) |
|
|
|
|
| def render_branch_graph(repo_path: str, next_commit_message: str) -> tuple[plt.Figure, str]: |
| repo = load_repo(repo_path) |
| commits = get_commit_graph(repo) |
| commits = tag_commits_with_branches(repo, commits) |
| current_branch = repo.active_branch.name if not repo.head.is_detached else "HEAD (detached)" |
| fig = build_figure(repo, commits, current_branch, next_commit_message) |
| info = build_info(repo, current_branch, next_commit_message) |
| return fig, info |
|
|
