677 lines
23 KiB
Python
677 lines
23 KiB
Python
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
import curses
|
||
import re
|
||
from typing import Iterable
|
||
|
||
from .remotes import (
|
||
RegistryError,
|
||
RemoteEntry,
|
||
RemoteQueryPool,
|
||
RemoteQueryResult,
|
||
RemoteRegistry,
|
||
)
|
||
|
||
|
||
SESSION_NAME_RE = re.compile(r"^[A-Za-z0-9._-]{1,64}$")
|
||
|
||
|
||
class SelectorError(RuntimeError):
|
||
"""The interactive selector could not be displayed."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CatalogWindow:
|
||
window_id: str
|
||
key: str
|
||
label: str
|
||
dead: bool
|
||
active: bool
|
||
kind: str | None = None
|
||
target: str | None = None
|
||
session: str | None = None
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SelectorResult:
|
||
kind: str
|
||
window_id: str | None = None
|
||
dead: bool = False
|
||
target: str | None = None
|
||
session: str | None = None
|
||
remote_keybindings: str = "local"
|
||
|
||
@classmethod
|
||
def window(cls, window_id: str, *, dead: bool) -> SelectorResult:
|
||
return cls(kind="window", window_id=window_id, dead=dead)
|
||
|
||
@classmethod
|
||
def close(cls, window_id: str) -> SelectorResult:
|
||
return cls(kind="close", window_id=window_id)
|
||
|
||
@classmethod
|
||
def local(cls, session: str) -> SelectorResult:
|
||
return cls(kind="local", session=session)
|
||
|
||
@classmethod
|
||
def remote(
|
||
cls,
|
||
target: str,
|
||
session: str,
|
||
remote_keybindings: str,
|
||
) -> SelectorResult:
|
||
return cls(
|
||
kind="remote",
|
||
target=target,
|
||
session=session,
|
||
remote_keybindings=remote_keybindings,
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PickerRow:
|
||
key: str
|
||
kind: str
|
||
main: str
|
||
detail: str = ""
|
||
selectable: bool = True
|
||
window_id: str | None = None
|
||
dead: bool = False
|
||
target: str | None = None
|
||
session: str | None = None
|
||
remote_keybindings: str = "local"
|
||
|
||
|
||
def build_main_rows(
|
||
windows: Iterable[CatalogWindow],
|
||
entries: Iterable[RemoteEntry],
|
||
results: dict[str, RemoteQueryResult],
|
||
loading_targets: set[str],
|
||
*,
|
||
local_sessions: Iterable[str] = (),
|
||
registry_error: str | None = None,
|
||
) -> list[PickerRow]:
|
||
window_list = list(windows)
|
||
entry_list = sorted(entries, key=lambda entry: entry.target)
|
||
registered = {entry.target for entry in entry_list}
|
||
rows: list[PickerRow] = []
|
||
|
||
standalone = [
|
||
window
|
||
for window in window_list
|
||
if window.kind == "local"
|
||
or window.target is None
|
||
or window.target not in registered
|
||
]
|
||
if standalone:
|
||
rows.append(PickerRow("heading:open", "heading", "已打开", selectable=False))
|
||
for window in standalone:
|
||
rows.append(_window_row(window))
|
||
|
||
open_local_sessions = {
|
||
window.session
|
||
for window in window_list
|
||
if window.kind == "local" and window.session is not None
|
||
}
|
||
available_local_sessions = sorted(
|
||
{
|
||
session
|
||
for session in local_sessions
|
||
if SESSION_NAME_RE.fullmatch(session) and session not in open_local_sessions
|
||
},
|
||
key=lambda name: (name != "default", name),
|
||
)
|
||
if available_local_sessions:
|
||
rows.append(PickerRow("heading:local", "heading", "本机", selectable=False))
|
||
rows.extend(
|
||
PickerRow(
|
||
f"local:{session}",
|
||
"local",
|
||
f"○ {session}",
|
||
"未在 Tab 中 · 单击打开",
|
||
session=session,
|
||
)
|
||
for session in available_local_sessions
|
||
)
|
||
|
||
active_window = next((window for window in window_list if window.active), None)
|
||
if active_window is not None:
|
||
can_close = len(window_list) > 1
|
||
rows.append(
|
||
PickerRow(
|
||
"close-active",
|
||
"close",
|
||
f"− 关闭当前 Tab：{active_window.label}",
|
||
"Herdr Session 保持运行" if can_close else "至少保留一个 Session Tab",
|
||
selectable=can_close,
|
||
window_id=active_window.window_id,
|
||
)
|
||
)
|
||
|
||
remote_windows = {
|
||
(window.target, window.session): window
|
||
for window in window_list
|
||
if window.kind == "remote"
|
||
and window.target is not None
|
||
and window.session is not None
|
||
}
|
||
for entry in entry_list:
|
||
rows.append(
|
||
PickerRow(
|
||
f"heading:{entry.target}",
|
||
"heading",
|
||
entry.target,
|
||
"远端",
|
||
selectable=False,
|
||
)
|
||
)
|
||
result = results.get(entry.target)
|
||
session_by_name = (
|
||
{
|
||
session.name: session
|
||
for session in result.sessions
|
||
if SESSION_NAME_RE.fullmatch(session.name)
|
||
}
|
||
if result is not None and result.error is None
|
||
else {}
|
||
)
|
||
for target, session_name in remote_windows:
|
||
if target == entry.target and session_name is not None:
|
||
session_by_name.setdefault(session_name, None)
|
||
|
||
for session_name in sorted(
|
||
session_by_name,
|
||
key=lambda name: (name != "default", name),
|
||
):
|
||
window = remote_windows.get((entry.target, session_name))
|
||
if window is not None:
|
||
rows.append(_window_row(window))
|
||
continue
|
||
session = session_by_name[session_name]
|
||
assert session is not None
|
||
rows.append(
|
||
PickerRow(
|
||
f"remote:{entry.target}:{session_name}",
|
||
"remote",
|
||
f"{'○' if session.running else '■'} {session_name}",
|
||
"远端运行 · 单击打开" if session.running else "已停止 · 单击启动",
|
||
target=entry.target,
|
||
session=session_name,
|
||
remote_keybindings=entry.remote_keybindings,
|
||
)
|
||
)
|
||
|
||
if entry.target in loading_targets:
|
||
rows.append(
|
||
PickerRow(
|
||
f"loading:{entry.target}",
|
||
"loading",
|
||
"… 正在查询 Sessions",
|
||
selectable=False,
|
||
)
|
||
)
|
||
elif result is not None and result.error is not None:
|
||
rows.append(
|
||
PickerRow(
|
||
f"retry:{entry.target}",
|
||
"retry",
|
||
"! 查询失败 · 单击重试",
|
||
result.error,
|
||
target=entry.target,
|
||
)
|
||
)
|
||
elif result is not None and not session_by_name:
|
||
rows.append(
|
||
PickerRow(
|
||
f"empty:{entry.target}",
|
||
"empty",
|
||
"没有已有 Session",
|
||
selectable=False,
|
||
)
|
||
)
|
||
rows.append(
|
||
PickerRow(
|
||
f"new:{entry.target}",
|
||
"new",
|
||
"+ 新建 Session…",
|
||
entry.target,
|
||
target=entry.target,
|
||
remote_keybindings=entry.remote_keybindings,
|
||
)
|
||
)
|
||
|
||
if registry_error is not None:
|
||
rows.append(
|
||
PickerRow(
|
||
"registry-error",
|
||
"error",
|
||
"! 远端配置不可用",
|
||
registry_error,
|
||
selectable=False,
|
||
)
|
||
)
|
||
else:
|
||
rows.extend(
|
||
[
|
||
PickerRow("add-remote", "add", "+ 添加远端…", "SSH target"),
|
||
PickerRow("manage-remotes", "manage", "⚙ 管理远端…"),
|
||
]
|
||
)
|
||
return rows
|
||
|
||
|
||
def _window_row(window: CatalogWindow) -> PickerRow:
|
||
symbol = "▶" if window.active else "×" if window.dead else "●"
|
||
detail = "当前" if window.active else "已断线 · 单击重连" if window.dead else "Hub 已打开"
|
||
return PickerRow(
|
||
f"window:{window.window_id}",
|
||
"window",
|
||
f"{symbol} {window.label}",
|
||
detail,
|
||
window_id=window.window_id,
|
||
dead=window.dead,
|
||
target=window.target,
|
||
session=window.session,
|
||
)
|
||
|
||
|
||
def filter_rows(rows: Iterable[PickerRow], query: str) -> list[PickerRow]:
|
||
if not query:
|
||
return list(rows)
|
||
needle = query.casefold()
|
||
return [
|
||
row
|
||
for row in rows
|
||
if needle in f"{row.main} {row.detail}".casefold()
|
||
]
|
||
|
||
|
||
class SessionSelector:
|
||
def __init__(
|
||
self,
|
||
windows: Iterable[CatalogWindow],
|
||
registry: RemoteRegistry,
|
||
ssh_bin: str,
|
||
local_sessions: Iterable[str] = (),
|
||
) -> None:
|
||
self.windows = list(windows)
|
||
self.local_sessions = list(local_sessions)
|
||
self.registry = registry
|
||
self.pool = RemoteQueryPool(ssh_bin)
|
||
self.entries: dict[str, RemoteEntry] = {}
|
||
self.results: dict[str, RemoteQueryResult] = {}
|
||
self.registry_error: str | None = None
|
||
self.message = ""
|
||
self.query = ""
|
||
self.selected_key: str | None = None
|
||
self.scroll = 0
|
||
self.view = "main"
|
||
self.input_value = ""
|
||
self.input_target: str | None = None
|
||
self._drawn_rows: dict[int, PickerRow] = {}
|
||
self._load_registry()
|
||
|
||
def _load_registry(self) -> None:
|
||
try:
|
||
entries = self.registry.load()
|
||
except RegistryError as error:
|
||
self.registry_error = str(error)
|
||
return
|
||
self.registry_error = None
|
||
self.entries = {entry.target: entry for entry in entries}
|
||
self.pool.enqueue_many(entries)
|
||
|
||
def run(self) -> SelectorResult | None:
|
||
try:
|
||
return curses.wrapper(self._curses_main)
|
||
except curses.error as error:
|
||
raise SelectorError(f"could not initialize the terminal selector: {error}") from error
|
||
finally:
|
||
self.pool.close()
|
||
|
||
def _curses_main(self, screen) -> SelectorResult | None: # type: ignore[no-untyped-def]
|
||
screen.keypad(True)
|
||
screen.timeout(100)
|
||
try:
|
||
curses.curs_set(0)
|
||
except curses.error:
|
||
pass
|
||
try:
|
||
curses.use_default_colors()
|
||
except curses.error:
|
||
pass
|
||
curses.mouseinterval(150)
|
||
curses.mousemask(curses.ALL_MOUSE_EVENTS)
|
||
|
||
while True:
|
||
for result in self.pool.poll():
|
||
self.results[result.entry.target] = result
|
||
rows = self._rows()
|
||
filtered = filter_rows(rows, self.query) if self.view == "main" else rows
|
||
self._normalize_selection(filtered)
|
||
self._draw(screen, filtered)
|
||
try:
|
||
key = screen.get_wch()
|
||
except curses.error:
|
||
continue
|
||
|
||
outcome = self._handle_key(key, filtered)
|
||
if isinstance(outcome, SelectorResult) or outcome is False:
|
||
return outcome if isinstance(outcome, SelectorResult) else None
|
||
|
||
def _rows(self) -> list[PickerRow]:
|
||
if self.view == "manage":
|
||
rows = [PickerRow("manage-heading", "heading", "管理远端", selectable=False)]
|
||
rows.extend(
|
||
PickerRow(
|
||
f"remove:{entry.target}",
|
||
"remove",
|
||
f"− {entry.target}",
|
||
"单击后确认移除",
|
||
target=entry.target,
|
||
)
|
||
for entry in sorted(self.entries.values(), key=lambda item: item.target)
|
||
)
|
||
if len(rows) == 1:
|
||
rows.append(PickerRow("no-remotes", "empty", "没有已登记远端", selectable=False))
|
||
rows.append(PickerRow("manage-back", "back", "← 返回 Sessions"))
|
||
return rows
|
||
if self.view == "confirm-remove":
|
||
target = self.input_target or ""
|
||
return [
|
||
PickerRow("confirm-heading", "heading", f"移除 {target}？", "不会关闭已有窗口或远端 server", selectable=False),
|
||
PickerRow("confirm-yes", "confirm-yes", "移除", target=target),
|
||
PickerRow("confirm-no", "back", "取消"),
|
||
]
|
||
if self.view in {"input-add", "input-new"}:
|
||
return []
|
||
return build_main_rows(
|
||
self.windows,
|
||
self.entries.values(),
|
||
self.results,
|
||
{target for target in self.entries if self.pool.is_loading(target)},
|
||
local_sessions=self.local_sessions,
|
||
registry_error=self.registry_error,
|
||
)
|
||
|
||
def _handle_key(self, key, rows: list[PickerRow]): # type: ignore[no-untyped-def]
|
||
if self.view in {"input-add", "input-new"}:
|
||
return self._handle_input_key(key)
|
||
if key == curses.KEY_MOUSE:
|
||
return self._handle_mouse(rows)
|
||
if key in {curses.KEY_UP}:
|
||
self._move(rows, -1)
|
||
return None
|
||
if key in {curses.KEY_DOWN}:
|
||
self._move(rows, 1)
|
||
return None
|
||
if key == curses.KEY_PPAGE:
|
||
self._move(rows, -8)
|
||
return None
|
||
if key == curses.KEY_NPAGE:
|
||
self._move(rows, 8)
|
||
return None
|
||
if key in {curses.KEY_ENTER, "\n", "\r"}:
|
||
return self._activate(self._selected_row(rows))
|
||
if key == "\x12" and self.view == "main":
|
||
self.results.clear()
|
||
self.pool.enqueue_many(self.entries.values())
|
||
self.message = "正在刷新所有远端"
|
||
return None
|
||
if key == "\x1b":
|
||
if self.view != "main":
|
||
self._show_main()
|
||
return None
|
||
if self.query:
|
||
self.query = ""
|
||
self.selected_key = None
|
||
return None
|
||
return False
|
||
if key in {curses.KEY_BACKSPACE, "\x7f", "\b"} and self.view == "main":
|
||
self.query = self.query[:-1]
|
||
self.selected_key = None
|
||
return None
|
||
if isinstance(key, str) and key.isprintable() and self.view == "main":
|
||
if len(self.query) < 128:
|
||
self.query += key
|
||
self.selected_key = None
|
||
return None
|
||
return None
|
||
|
||
def _handle_input_key(self, key): # type: ignore[no-untyped-def]
|
||
if key == "\x1b":
|
||
self._show_main()
|
||
return None
|
||
if key in {curses.KEY_ENTER, "\n", "\r"}:
|
||
value = self.input_value.strip()
|
||
if self.view == "input-add":
|
||
return self._finish_add(value)
|
||
return self._finish_new(value)
|
||
if key in {curses.KEY_BACKSPACE, "\x7f", "\b"}:
|
||
self.input_value = self.input_value[:-1]
|
||
return None
|
||
if isinstance(key, str) and key.isprintable() and len(self.input_value) < 512:
|
||
self.input_value += key
|
||
return None
|
||
|
||
def _handle_mouse(self, rows: list[PickerRow]):
|
||
try:
|
||
_, _, y, _, state = curses.getmouse()
|
||
except curses.error:
|
||
return None
|
||
if state & getattr(curses, "BUTTON4_PRESSED", 0):
|
||
self._move(rows, -3)
|
||
return None
|
||
if state & getattr(curses, "BUTTON5_PRESSED", 0):
|
||
self._move(rows, 3)
|
||
return None
|
||
clicked = (
|
||
getattr(curses, "BUTTON1_CLICKED", 0)
|
||
| getattr(curses, "BUTTON1_DOUBLE_CLICKED", 0)
|
||
)
|
||
if state & clicked:
|
||
row = self._drawn_rows.get(y)
|
||
if row is not None and row.selectable:
|
||
self.selected_key = row.key
|
||
return self._activate(row)
|
||
return None
|
||
|
||
def _activate(self, row: PickerRow | None):
|
||
if row is None:
|
||
return None
|
||
if row.kind == "window":
|
||
assert row.window_id is not None
|
||
return SelectorResult.window(row.window_id, dead=row.dead)
|
||
if row.kind == "close":
|
||
assert row.window_id is not None
|
||
return SelectorResult.close(row.window_id)
|
||
if row.kind == "local":
|
||
assert row.session is not None
|
||
return SelectorResult.local(row.session)
|
||
if row.kind == "remote":
|
||
assert row.target is not None and row.session is not None
|
||
return SelectorResult.remote(
|
||
row.target,
|
||
row.session,
|
||
row.remote_keybindings,
|
||
)
|
||
if row.kind == "retry":
|
||
assert row.target is not None
|
||
self.results.pop(row.target, None)
|
||
self.pool.enqueue(self.entries[row.target])
|
||
self.message = f"正在重试 {row.target}"
|
||
elif row.kind == "new":
|
||
self.view = "input-new"
|
||
self.input_target = row.target
|
||
self.input_value = ""
|
||
self.message = ""
|
||
elif row.kind == "add":
|
||
self.view = "input-add"
|
||
self.input_target = None
|
||
self.input_value = ""
|
||
self.message = ""
|
||
elif row.kind == "manage":
|
||
self.view = "manage"
|
||
self.query = ""
|
||
self.selected_key = None
|
||
elif row.kind == "remove":
|
||
self.view = "confirm-remove"
|
||
self.input_target = row.target
|
||
self.selected_key = "confirm-no"
|
||
elif row.kind == "confirm-yes":
|
||
self._finish_remove(row.target or "")
|
||
elif row.kind == "back":
|
||
self._show_main()
|
||
return None
|
||
|
||
def _finish_add(self, target: str):
|
||
try:
|
||
entry = RemoteEntry(target)
|
||
self.registry.upsert(entry)
|
||
except (ValueError, RegistryError) as error:
|
||
self.message = str(error)
|
||
return None
|
||
self.entries[target] = entry
|
||
self.results.pop(target, None)
|
||
self.pool.enqueue(entry)
|
||
self._show_main()
|
||
self.message = f"已添加 {target}，正在查询"
|
||
return None
|
||
|
||
def _finish_new(self, session: str):
|
||
if not SESSION_NAME_RE.fullmatch(session):
|
||
self.message = "Session 名称需为 1-64 位字母、数字、. _ -"
|
||
return None
|
||
target = self.input_target
|
||
assert target is not None
|
||
entry = self.entries[target]
|
||
return SelectorResult.remote(target, session, entry.remote_keybindings)
|
||
|
||
def _finish_remove(self, target: str) -> None:
|
||
try:
|
||
self.registry.remove(target)
|
||
except RegistryError as error:
|
||
self.message = str(error)
|
||
self._show_main()
|
||
return
|
||
self.pool.cancel(target)
|
||
self.entries.pop(target, None)
|
||
self.results.pop(target, None)
|
||
self._show_main()
|
||
self.message = f"已移除 {target}；已有 window 保持不变"
|
||
|
||
def _show_main(self) -> None:
|
||
self.view = "main"
|
||
self.input_target = None
|
||
self.input_value = ""
|
||
self.selected_key = None
|
||
self.scroll = 0
|
||
|
||
def _normalize_selection(self, rows: list[PickerRow]) -> None:
|
||
selectable = [row for row in rows if row.selectable]
|
||
if not selectable:
|
||
self.selected_key = None
|
||
self.scroll = 0
|
||
return
|
||
if self.selected_key not in {row.key for row in selectable}:
|
||
self.selected_key = selectable[0].key
|
||
|
||
def _selected_row(self, rows: list[PickerRow]) -> PickerRow | None:
|
||
return next((row for row in rows if row.key == self.selected_key), None)
|
||
|
||
def _move(self, rows: list[PickerRow], delta: int) -> None:
|
||
selectable = [row for row in rows if row.selectable]
|
||
if not selectable:
|
||
return
|
||
current = next(
|
||
(index for index, row in enumerate(selectable) if row.key == self.selected_key),
|
||
0,
|
||
)
|
||
next_index = max(0, min(len(selectable) - 1, current + delta))
|
||
self.selected_key = selectable[next_index].key
|
||
|
||
def _draw(self, screen, rows: list[PickerRow]) -> None: # type: ignore[no-untyped-def]
|
||
screen.erase()
|
||
height, width = screen.getmaxyx()
|
||
self._drawn_rows = {}
|
||
if height < 8 or width < 28:
|
||
self._safe_add(screen, 0, 0, "终端太小，请旋转屏幕", curses.A_BOLD)
|
||
screen.refresh()
|
||
return
|
||
|
||
if self.view in {"input-add", "input-new"}:
|
||
self._draw_input(screen, height, width)
|
||
return
|
||
try:
|
||
curses.curs_set(0)
|
||
except curses.error:
|
||
pass
|
||
|
||
title = "Sessions" if self.view == "main" else "管理远端"
|
||
self._safe_add(screen, 1, 2, title, curses.A_BOLD)
|
||
if self.view == "main":
|
||
self._safe_add(screen, 2, 2, f"搜索: {self.query}", curses.A_DIM)
|
||
if self.message:
|
||
self._safe_add(screen, 3, 2, self.message, curses.A_DIM)
|
||
|
||
list_top = 4
|
||
list_height = max(1, height - list_top - 2)
|
||
selected_index = next(
|
||
(index for index, row in enumerate(rows) if row.key == self.selected_key),
|
||
0,
|
||
)
|
||
if selected_index < self.scroll:
|
||
self.scroll = selected_index
|
||
elif selected_index >= self.scroll + list_height:
|
||
self.scroll = selected_index - list_height + 1
|
||
self.scroll = max(0, min(self.scroll, max(0, len(rows) - list_height)))
|
||
|
||
for offset, row in enumerate(rows[self.scroll : self.scroll + list_height]):
|
||
y = list_top + offset
|
||
self._drawn_rows[y] = row
|
||
selected = row.key == self.selected_key
|
||
attribute = curses.A_REVERSE if selected else curses.A_NORMAL
|
||
if row.kind == "heading":
|
||
attribute |= curses.A_BOLD
|
||
elif not row.selectable:
|
||
attribute |= curses.A_DIM
|
||
detail_room = max(0, width // 2 - 2)
|
||
main_room = max(8, width - detail_room - 6)
|
||
main = row.main[:main_room]
|
||
detail = row.detail[:detail_room]
|
||
line = f" {main:<{main_room}} {detail}"
|
||
self._safe_add(screen, y, 0, line, attribute)
|
||
|
||
footer = "单击/Enter 打开 · 输入搜索 · Ctrl+R 刷新 · Esc 关闭"
|
||
if self.view != "main":
|
||
footer = "单击/Enter 选择 · Esc 返回"
|
||
self._safe_add(screen, height - 1, 1, footer, curses.A_DIM)
|
||
screen.refresh()
|
||
|
||
def _draw_input(self, screen, height: int, width: int) -> None: # type: ignore[no-untyped-def]
|
||
is_add = self.view == "input-add"
|
||
title = "添加远端" if is_add else f"在 {self.input_target} 新建 Session"
|
||
prompt = "SSH target" if is_add else "Session 名称"
|
||
self._safe_add(screen, 1, 2, title, curses.A_BOLD)
|
||
self._safe_add(screen, 3, 2, f"{prompt}: {self.input_value}")
|
||
if self.message:
|
||
self._safe_add(screen, 5, 2, self.message, curses.A_DIM)
|
||
self._safe_add(screen, height - 1, 1, "Enter 确认 · Esc 返回", curses.A_DIM)
|
||
try:
|
||
curses.curs_set(1)
|
||
screen.move(3, min(width - 2, len(prompt) + len(self.input_value) + 4))
|
||
except curses.error:
|
||
pass
|
||
screen.refresh()
|
||
|
||
@staticmethod
|
||
def _safe_add(screen, y: int, x: int, text: str, attribute: int = curses.A_NORMAL) -> None: # type: ignore[no-untyped-def]
|
||
height, width = screen.getmaxyx()
|
||
if y < 0 or y >= height or x >= width:
|
||
return
|
||
try:
|
||
screen.addnstr(y, x, text, max(0, width - x - 1), attribute)
|
||
except curses.error:
|
||
pass
|