Files
lightweight-cobot/cobot/commands/local_setup.py
T

389 lines
13 KiB
Python

from __future__ import annotations
import argparse
import os
import shutil
import subprocess
import tempfile
import threading
import urllib.request
from pathlib import Path
from typing import Callable, List, Optional
from textual.app import App
from cobot.tui import SCREEN_CSS, LogScreen, PickScreen
from cobot.commands.docker_setup import run as _docker_setup
_PROJECT_DIR = Path(__file__).parent.parent.parent
_ROS_KEYRING = Path("/usr/share/keyrings/ros-archive-keyring.gpg")
_ROS_SOURCES = Path("/etc/apt/sources.list.d/ros2.list")
_ROS_KEY_URL = "https://raw.githubusercontent.com/ros/rosdistro/master/ros.key"
_APT_ENV = {**os.environ, "DEBIAN_FRONTEND": "noninteractive"}
def _detect_ubuntu_2404() -> bool:
path = Path("/etc/os-release")
if not path.exists():
return False
info: dict[str, str] = {}
for line in path.read_text().splitlines():
if "=" in line:
k, _, v = line.partition("=")
info[k.strip()] = v.strip().strip('"')
return info.get("ID") == "ubuntu" and info.get("VERSION_ID") == "24.04"
def _detect_ros2_jazzy() -> bool:
return Path("/opt/ros/jazzy").is_dir()
Write = Callable[[str], None]
def _run_quiet(cmd: List[str], write: Write | None = None, env: dict | None = None, cwd=None) -> None:
result = subprocess.run(
cmd, capture_output=True, text=True,
env=env or os.environ, cwd=cwd,
)
if result.returncode != 0:
if write:
for line in (result.stdout + result.stderr).splitlines():
if line.strip():
write(line)
raise RuntimeError(f"Command failed: {cmd[0]}")
def _run_logged(cmd: List[str], write: Write, env: dict | None = None, cwd=None) -> None:
proc = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
env=env or os.environ,
cwd=cwd,
)
for line in proc.stdout:
s = line.rstrip()
if s:
write(s)
proc.wait()
if proc.returncode != 0:
raise RuntimeError(f"Command failed: {cmd[0]}")
def _run_apt_with_progress(
cmd: List[str],
write: Write,
on_progress: Callable[[float], None],
env: dict | None = None,
) -> None:
"""Run an apt command and feed real percentage from APT::Status-Fd to on_progress(0-100)."""
r_fd, w_fd = os.pipe()
try:
proc = subprocess.Popen(
cmd + [f"-o", f"APT::Status-Fd={w_fd}"],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
env=env or os.environ,
pass_fds=(w_fd,),
)
finally:
os.close(w_fd)
def _read_status() -> None:
with os.fdopen(r_fd, "r") as f:
for line in f:
# Format: dlstatus:N:PCT:MSG or pmstatus:NAME:PCT:MSG
parts = line.strip().split(":", 3)
if len(parts) >= 3:
try:
on_progress(float(parts[2]))
except ValueError:
pass
t = threading.Thread(target=_read_status, daemon=True)
t.start()
for line in proc.stdout:
s = line.rstrip()
if s:
write(s)
proc.wait()
t.join()
if proc.returncode != 0:
raise RuntimeError(f"Command failed: {cmd[0]}")
# ---------------------------------------------------------------------------
# Installation steps
# ---------------------------------------------------------------------------
def _setup_locale(write: Write) -> None:
write("[cyan][*][/cyan] Checking locale...")
if "UTF-8" in subprocess.run(["locale"], capture_output=True, text=True).stdout:
write("[green][ok][/green] UTF-8 locale active")
return
write("[cyan][*][/cyan] Configuring UTF-8 locale...")
_run_quiet(["sudo", "apt-get", "update", "-qq"], write)
_run_quiet(["sudo", "apt-get", "install", "-y", "--no-install-recommends", "locales"], write, _APT_ENV)
_run_quiet(["sudo", "locale-gen", "en_US.UTF-8"], write)
_run_quiet(["sudo", "update-locale", "LC_ALL=en_US.UTF-8", "LANG=en_US.UTF-8"], write)
write("[green][ok][/green] Locale configured")
def _add_ros2_repo(write: Write, on_progress: Optional[Callable[[float], None]] = None) -> None:
def _prog(p: float) -> None:
if on_progress:
on_progress(p)
write("[cyan][*][/cyan] Adding ROS2 apt repository...")
# Best-effort update before installing prereqs
subprocess.run(["sudo", "apt-get", "update", "-qq"], capture_output=True)
_prog(15)
_run_quiet(
["sudo", "apt-get", "install", "-y", "--no-install-recommends",
"software-properties-common", "curl", "gnupg"],
write, _APT_ENV,
)
_prog(35)
_run_quiet(["sudo", "add-apt-repository", "-y", "universe"], write)
_prog(45)
write("[cyan][*][/cyan] Downloading ROS2 signing key...")
with tempfile.NamedTemporaryFile(delete=False, suffix=".key") as tmp:
tmp_path = tmp.name
try:
urllib.request.urlretrieve(_ROS_KEY_URL, tmp_path)
_prog(60)
_run_quiet(["sudo", "gpg", "--dearmor", "--yes", "-o", str(_ROS_KEYRING), tmp_path])
finally:
os.unlink(tmp_path)
write("[green][ok][/green] Signing key installed")
_prog(65)
arch = subprocess.check_output(["dpkg", "--print-architecture"], text=True).strip()
codename = subprocess.check_output(
["bash", "-c", ". /etc/os-release && echo $UBUNTU_CODENAME"], text=True
).strip()
sources_line = (
f"deb [arch={arch} signed-by={_ROS_KEYRING}] "
f"http://packages.ros.org/ros2/ubuntu {codename} main\n"
)
proc = subprocess.run(
["sudo", "tee", str(_ROS_SOURCES)],
input=sources_line, capture_output=True, text=True,
)
if proc.returncode != 0:
raise RuntimeError(f"Failed to write {_ROS_SOURCES}")
_prog(70)
write("[cyan][*][/cyan] Updating apt cache...")
_run_apt_with_progress(
["sudo", "apt-get", "update", "-q"],
write,
lambda p: _prog(70 + p * 0.30),
_APT_ENV,
)
write("[green][ok][/green] ROS2 repository ready")
_prog(100)
def _install_ros2_jazzy(write: Write, on_progress: Optional[Callable[[float], None]] = None) -> None:
write("[cyan][*][/cyan] Installing ros-jazzy-desktop and ros-dev-tools...")
_run_apt_with_progress(
["sudo", "apt-get", "install", "-y", "ros-jazzy-desktop", "ros-dev-tools"],
write,
on_progress or (lambda _: None),
_APT_ENV,
)
write("[green][ok][/green] ROS2 Jazzy Desktop installed")
def _install_colcon(write: Write) -> None:
if shutil.which("colcon"):
write("[green][ok][/green] colcon already available")
return
write("[cyan][*][/cyan] Installing colcon...")
_run_quiet(
["sudo", "apt-get", "install", "-y", "--no-install-recommends",
"python3-colcon-common-extensions"],
write, _APT_ENV,
)
write("[green][ok][/green] colcon installed")
def _setup_shell_rc(write: Write) -> None:
shell_name = Path(os.environ.get("SHELL", "/bin/bash")).name
rc = Path.home() / (".zshrc" if shell_name == "zsh" else ".bashrc")
source_line = "source /opt/ros/jazzy/setup.bash"
if rc.exists() and source_line in rc.read_text():
write(f"[green][ok][/green] ROS2 setup already in {rc.name}")
return
with rc.open("a") as f:
f.write(f"\n# ROS2 Jazzy\n{source_line}\n")
write(f"[green][ok][/green] Added ROS2 setup to ~/{rc.name}")
# ---------------------------------------------------------------------------
# Background tasks (run inside LogScreen worker)
# ---------------------------------------------------------------------------
def _task_install_jazzy(screen: LogScreen) -> None:
try:
# Step 1 — locale (0 → 5 %)
screen.set_progress(0, "Setting up locale...")
screen.write("[bold]Step 1 / 5 — Locale[/bold]")
_setup_locale(screen.write)
# Step 2 — ROS2 repo (5 → 20 %)
screen.set_progress(5, "Adding ROS2 repository...")
screen.write("\n[bold]Step 2 / 5 — ROS2 repository[/bold]")
_add_ros2_repo(
screen.write,
on_progress=lambda p: screen.set_progress(5 + p * 0.15),
)
# Step 3 — ROS2 Jazzy (20 → 85 %)
screen.set_progress(20, "Installing ROS2 Jazzy Desktop...")
screen.write("\n[bold]Step 3 / 5 — ROS2 Jazzy Desktop[/bold]")
_install_ros2_jazzy(
screen.write,
on_progress=lambda p: screen.set_progress(20 + p * 0.65),
)
# Step 4 — colcon (85 → 92 %)
screen.set_progress(85, "Installing colcon...")
screen.write("\n[bold]Step 4 / 5 — colcon[/bold]")
_install_colcon(screen.write)
# Step 5 — shell rc (92 → 100 %)
screen.set_progress(92, "Configuring shell...")
screen.write("\n[bold]Step 5 / 5 — Shell configuration[/bold]")
_setup_shell_rc(screen.write)
screen.set_progress(100, "Done")
screen.write(
"\nRestart the terminal, then run [bold]cobot local-setup[/bold] again to build."
)
screen.finish(True)
except Exception as exc:
screen.write(f"\n[red]Error:[/red] {exc}")
screen.finish(False)
def _task_build(screen: LogScreen) -> None:
try:
if not shutil.which("colcon"):
screen.write("[red]colcon not found.[/red]")
screen.write("Source ROS2 first: [bold]source /opt/ros/jazzy/setup.bash[/bold]")
screen.finish(False)
return
# Count packages so we can show X/total progress
list_result = subprocess.run(
["colcon", "list"], capture_output=True, text=True, cwd=_PROJECT_DIR,
)
total = max(len([l for l in list_result.stdout.splitlines() if l.strip()]), 1)
screen.write(f"[bold]Building {total} package(s) with colcon[/bold]\n")
screen.set_progress(0, f"0 / {total} packages done")
built = 0
def _track(line: str) -> None:
nonlocal built
screen.write(line)
if "Finished <<<" in line or "Failed <<<" in line:
built += 1
screen.set_progress(built / total * 100, f"{built} / {total} packages done")
_run_logged(["colcon", "build", "--symlink-install"], _track, cwd=_PROJECT_DIR)
screen.set_progress(100, "Build complete")
screen.write("\nActivate workspace: [bold]source install/setup.bash[/bold]")
screen.finish(True)
except Exception as exc:
screen.write(f"\n[red]Error:[/red] {exc}")
screen.finish(False)
# ---------------------------------------------------------------------------
# Textual apps
# ---------------------------------------------------------------------------
class _InstallJazzyApp(App[None]):
CSS = SCREEN_CSS
def on_mount(self) -> None:
self.push_screen(
PickScreen(
"ROS2 not found",
"ROS2 Jazzy is not installed. Install it now?",
["Yes, install ROS2 Jazzy", "No, skip"],
"Yes, install ROS2 Jazzy",
),
self._on_choice,
)
def _on_choice(self, choice: Optional[str]) -> None:
if choice is None or choice.startswith("No"):
self.exit()
return
self.push_screen(
LogScreen("Installing ROS2 Jazzy", _task_install_jazzy, show_progress=True),
lambda _: self.exit(),
)
class _BuildApp(App[None]):
CSS = SCREEN_CSS
def on_mount(self) -> None:
self.push_screen(
LogScreen("Building project", _task_build, show_progress=True),
lambda _: self.exit(),
)
class _DockerPromptApp(App[bool]):
CSS = SCREEN_CSS
def on_mount(self) -> None:
self.push_screen(
PickScreen(
"Unsupported OS",
"Ubuntu 24.04 not detected. Build a Docker image for development?",
["Yes, run docker-setup", "No, exit"],
"Yes, run docker-setup",
),
lambda v: self.exit(v is not None and v.startswith("Yes")),
)
# ---------------------------------------------------------------------------
# CLI registration
# ---------------------------------------------------------------------------
def register(subparsers: argparse._SubParsersAction) -> None:
p = subparsers.add_parser(
"local-setup",
help="Install ROS2 Jazzy natively and build the project with colcon",
)
p.set_defaults(func=run)
def run(args: argparse.Namespace) -> None:
if not _detect_ubuntu_2404():
if _DockerPromptApp().run():
_docker_setup(args)
return
if not _detect_ros2_jazzy():
_InstallJazzyApp().run()
return
_BuildApp().run()