from __future__ import annotations import argparse import os import re import shutil import subprocess import sys from dataclasses import dataclass from pathlib import Path from typing import Callable, List, Optional from textual.app import App from cobot.tui import SCREEN_CSS, InputScreen, LogScreen, PickScreen _PROJECT_DIR = Path(__file__).parent.parent.parent _DOCKER_DIR = _PROJECT_DIR / "docker" # Default Docker Hub repository and local image prefix used when building locally. # Репозиторий Docker Hub по умолчанию и локальный префикс образов при локальной сборке. _DEFAULT_HUB_REPO = "evilfisru/lwc" _DEFAULT_PREFIX = "lwc-local" # The images must be built in this order because each one is based on the previous. # Образы должны собираться в этом порядке, потому что каждый основан на предыдущем. _CONTROLLER_CHAIN = ["ros-core", "ros-base", "ros-iiwa7"] _WEBOTS_CHAIN = ["ros-core", "ros-base", "ros-iiwa7-webots"] # Maps each image to the image it is built FROM. None means it starts from scratch (base Ubuntu). # Сопоставляет каждый образ с тем, на основе которого он собирается. None - начинает с нуля (базовый Ubuntu). _IMAGE_PARENT: dict[str, str | None] = { "ros-core": None, "ros-base": "ros-core", "ros-iiwa7": "ros-base", "ros-iiwa7-webots": "ros-base", } # These images need the full project source as Docker build context because they copy source files. # Эти образы требуют полный исходный код проекта как контекст сборки, потому что копируют файлы. _NEEDS_PROJECT_CTX = {"ros-iiwa7", "ros-iiwa7-webots"} Write = Callable[[str], None] # All the choices the user makes in the wizard are stored here before we start the actual build. # Все выборы пользователя в мастере хранятся здесь перед началом фактической сборки. @dataclass class _Config: ros_version: str variant: str source: str build_type: str image_prefix: str hub_repo: str # Build one Docker image and stream its output to the log. # Progress is tracked by parsing "Step X/Y" lines that Docker prints during the build. # Собирает один Docker-образ и транслирует его вывод в лог. # Прогресс отслеживается по строкам "Step X/Y", которые Docker печатает во время сборки. def _build_image( name: str, tag: str, dockerfile: Path, ctx: Path, write: Write, on_progress: Optional[Callable[[float], None]] = None, parent_tag: Optional[str] = None, build_type: str = "release", register_proc: Optional[Callable] = None, ) -> bool: """Build a single Docker image from a Dockerfile and stream its output line by line. Returns True on success, False if the build failed or was cancelled. Собирает один Docker-образ из Dockerfile и транслирует вывод построчно. Возвращает True при успехе, False если сборка завершилась ошибкой или была отменена. """ write(f"[cyan][*][/cyan] Building [bold]{name}[/bold]...") # DOCKER_BUILDKIT=0 gives us "Step X/Y" lines in the output which we parse for progress. # DOCKER_BUILDKIT=0 даёт нам строки "Step X/Y" в выводе, которые мы парсим для прогресса. env = {**os.environ, "DOCKER_BUILDKIT": "0"} cmd = [ "docker", "build", "-t", tag, "-f", str(dockerfile), "--build-arg", f"BUILD_TYPE={build_type}", ] if parent_tag: cmd += ["--build-arg", f"IMAGE={parent_tag}"] cmd.append(str(ctx)) proc = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, env=env, ) if register_proc: register_proc(proc) for line in proc.stdout: s = line.rstrip() if s: write(s) if on_progress: m = re.match(r"Step (\d+)/(\d+) :", line) if m: step, total = int(m.group(1)), int(m.group(2)) on_progress(step / total * 100) proc.wait() if proc.returncode in (-9, -15): return False if proc.returncode == 0: write(f"[green][ok][/green] {name}") return True write(f"[red]Build failed:[/red] {tag}") return False # Pull a Docker image from Hub and track progress by counting downloaded layers. # Скачиваем Docker-образ с Hub и отслеживаем прогресс по количеством скачанных слоёв. def _pull_image( name: str, tag: str, write: Write, on_progress: Optional[Callable[[float], None]] = None, register_proc: Optional[Callable] = None, ) -> bool: """Pull a Docker image from Docker Hub and report layer-by-layer progress. Returns True on success, False if the pull failed or was cancelled. Скачивает Docker-образ с Docker Hub и сообщает о прогрессе по слоям. Возвращает True при успехе, False если скачивание завершилось ошибкой или было отменено. """ write(f"[cyan][*][/cyan] Pulling [bold]{name}[/bold] ({tag})...") if on_progress: on_progress(5) proc = subprocess.Popen( ["docker", "pull", tag], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, ) if register_proc: register_proc(proc) layers_total = 0 layers_done = 0 for line in proc.stdout: s = line.rstrip() if s: write(s) # Count layers as they appear and mark them done when Docker confirms they are pulled. # Считаем слои по мере их появления и отмечаем завершёнными когда Docker подтверждает скачивание. if "Pulling fs layer" in line or "Waiting" in line: layers_total += 1 elif "Pull complete" in line or "Already exists" in line: layers_done += 1 if on_progress and layers_total > 0: on_progress(5 + layers_done / layers_total * 90) proc.wait() if proc.returncode in (-9, -15): return False if proc.returncode == 0: write(f"[green][ok][/green] {name}") if on_progress: on_progress(100) return True write(f"[red]Pull failed:[/red] {tag}") return False # The actual work - either build or pull all images depending on what the user chose. # Each image gets its own slice of the progress bar so the overall bar advances smoothly. # Основная работа - собираем или скачиваем все образы в зависимости от выбора пользователя. # Каждый образ получает свой кусок прогресс-бара, чтобы общий бар двигался равномерно. def _task_execute(screen: LogScreen, cfg: _Config) -> None: """Worker function that runs inside LogScreen. Builds or pulls all images in the chain defined by the user's choices and updates the progress bar after each image. Рабочая функция, выполняемая внутри LogScreen. Собирает или скачивает все образы из цепочки согласно выбору пользователя и обновляет прогресс-бар после каждого образа. """ try: chain = _WEBOTS_CHAIN if cfg.variant == "webots" else _CONTROLLER_CHAIN n = len(chain) if cfg.source == "build": screen.write( f"[bold]Building {n} image(s) — " f"ROS {cfg.ros_version} — {cfg.build_type}[/bold]\n" ) for i, name in enumerate(chain): if screen.is_stopped(): return lo = i / n * 100 hi = (i + 1) / n * 100 screen.set_progress(lo, f"Image {i + 1}/{n}: building {name}...") tag = f"{cfg.image_prefix}:{name}-{cfg.ros_version}" dockerfile = _DOCKER_DIR / cfg.ros_version / name / "Dockerfile" if not dockerfile.exists(): screen.write(f"[red]Dockerfile not found:[/red] {dockerfile}") if not screen.is_stopped(): screen.finish(False) return ctx = _PROJECT_DIR if name in _NEEDS_PROJECT_CTX else dockerfile.parent parent_name = _IMAGE_PARENT.get(name) parent_tag = ( f"{cfg.image_prefix}:{parent_name}-{cfg.ros_version}" if parent_name else None ) ok = _build_image( name, tag, dockerfile, ctx, screen.write, on_progress=lambda p, lo=lo, hi=hi: screen.set_progress( lo + p * (hi - lo) / 100, f"Image {i + 1}/{n}: building {name}..." ), parent_tag=parent_tag, build_type=cfg.build_type, register_proc=screen.set_proc, ) if screen.is_stopped(): return if not ok: screen.finish(False) return screen.set_progress(hi) if not screen.is_stopped(): screen.set_progress(100, "All images built") screen.write( f"\n[green]Done.[/green] " f"Images tagged [bold]{cfg.image_prefix}:-{cfg.ros_version}[/bold]." ) else: short = "webots" if cfg.variant == "webots" else "iiwa" suffix = "-dev" if cfg.build_type == "dev" else "" full_ref = f"{cfg.hub_repo}:{short}-{cfg.ros_version}{suffix}" screen.write( f"[bold]Pulling from {cfg.hub_repo} — " f"ROS {cfg.ros_version} — {cfg.build_type}[/bold]\n" ) screen.set_progress(0, f"Pulling {full_ref}...") ok = _pull_image( short, full_ref, screen.write, on_progress=lambda p: screen.set_progress(p, f"Pulling {full_ref}..."), register_proc=screen.set_proc, ) if screen.is_stopped(): return if not ok: screen.finish(False) return screen.set_progress(100, "Pull complete") screen.write(f"\n[green]Done.[/green] Image ready: [bold]{full_ref}[/bold].") if not screen.is_stopped(): screen.finish(True) except Exception as exc: if not screen.is_stopped(): screen.write(f"\n[red]Error:[/red] {exc}") screen.finish(False) # Scan the docker/ directory for subdirectories named after ROS versions (e.g. jazzy). # If nothing is found we fall back to "jazzy" so the wizard still works. # Сканируем директорию docker/ на наличие поддиректорий с именами версий ROS (например jazzy). # Если ничего не найдено, используем "jazzy" по умолчанию, чтобы мастер всё равно работал. def _discover_versions() -> List[str]: """Return a sorted list of ROS versions found in docker/. Jazzy is placed first. Falls back to ["jazzy"] if the directory does not exist or is empty. Возвращает отсортированный список версий ROS найденных в docker/. Jazzy идёт первым. Возвращает ["jazzy"] если директория не существует или пуста. """ if not _DOCKER_DIR.exists(): return ["jazzy"] dirs = sorted(d.name for d in _DOCKER_DIR.iterdir() if d.is_dir()) # Put jazzy first so it is the pre-selected default in the wizard. # Ставим jazzy первым, чтобы он был предвыбранным по умолчанию в мастере. if "jazzy" in dirs: dirs = ["jazzy"] + [d for d in dirs if d != "jazzy"] return dirs or ["jazzy"] # Multi-step wizard that collects all build options before starting the actual image build. # Многошаговый мастер, который собирает все параметры сборки перед запуском фактической сборки образа. class _Wizard(App[None]): """Five-step wizard: ROS version -> source (pull/build) -> variant -> build type -> repo/prefix. Collects all options, then hands off to LogScreen which runs _task_execute. Пятишаговый мастер: версия ROS -> источник (pull/build) -> вариант -> тип сборки -> репо/префикс. Собирает все параметры, затем передаёт управление LogScreen который запускает _task_execute. """ CSS = SCREEN_CSS def __init__(self, versions: List[str], default_version: str = "jazzy"): super().__init__() self.versions = versions self.default_version = default_version self._state: dict = {} def on_mount(self) -> None: self._ask_version() def _ask_version(self) -> None: self.push_screen( PickScreen("Step 1 of 5", "Select ROS version:", self.versions, self.default_version), self._got_version, ) def _got_version(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["ros_version"] = v self.push_screen( PickScreen( "Step 2 of 5", "Source:", ["Pull from Docker Hub", "Build locally"], "Pull from Docker Hub", ), self._got_source, ) def _got_source(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["source"] = "build" if v == "Build locally" else "pull" self.push_screen( PickScreen( "Step 3 of 5", "What to install:", [ "Controller only — ros-core, ros-base, ros-iiwa7", "Controller with Webots — ros-core, ros-base, ros-iiwa7-webots", ], "Controller only — ros-core, ros-base, ros-iiwa7", ), self._got_variant, ) def _got_variant(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["variant"] = "webots" if v.startswith("Controller with Webots") else "controller" self.push_screen( PickScreen("Step 4 of 5", "Build type:", ["release", "dev"], "release"), self._got_build_type, ) def _got_build_type(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["build_type"] = v or "release" # Pull needs a Hub repo name, build needs a local image prefix. # Для pull нужно имя репозитория на Hub, для build - локальный префикс образов. if self._state["source"] == "pull": self.push_screen( InputScreen("Step 5 of 5", "Docker Hub repository:", _DEFAULT_HUB_REPO), self._got_hub_repo, ) else: self.push_screen( InputScreen("Step 5 of 5", "Image prefix:", _DEFAULT_PREFIX), self._got_prefix, ) def _got_hub_repo(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["hub_repo"] = v self._finish() def _got_prefix(self, v: Optional[str]) -> None: if v is None: self.exit() return self._state["image_prefix"] = v self._finish() def _finish(self) -> None: # Assemble the config and hand it off to the log screen that does the actual work. # Собираем конфиг и передаём его экрану лога, который выполняет фактическую работу. s = self._state cfg = _Config( ros_version=s["ros_version"], variant=s["variant"], source=s["source"], build_type=s["build_type"], image_prefix=s.get("image_prefix", _DEFAULT_PREFIX), hub_repo=s.get("hub_repo", _DEFAULT_HUB_REPO), ) title = ( f"Building Docker images — ROS {cfg.ros_version}" if cfg.source == "build" else f"Pulling Docker image — ROS {cfg.ros_version}" ) self.push_screen( LogScreen(title, lambda screen: _task_execute(screen, cfg), show_progress=True), lambda _: self.exit(), ) def register(subparsers: argparse._SubParsersAction) -> None: p = subparsers.add_parser("docker-setup", help="Build or pull Docker images for KUKA iiwa7") p.set_defaults(func=run) def run(args: argparse.Namespace) -> None: if not shutil.which("docker"): from rich.console import Console Console().print("[red]Error:[/red] Docker is not installed or not on PATH.") sys.exit(1) versions = _discover_versions() default = "jazzy" if "jazzy" in versions else versions[0] _Wizard(versions=versions, default_version=default).run()