updating and refactoring codebase for improved performance and maintainability
This commit is contained in:
+80
-56
@@ -1,5 +1,7 @@
|
||||
"""Data models for the pipeline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import hashlib
|
||||
import inspect
|
||||
@@ -16,23 +18,9 @@ from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional, Protocol, TextIO
|
||||
|
||||
from rich.console import Console
|
||||
from rich.console import ConsoleOptions
|
||||
from rich.console import Group
|
||||
from rich.live import Live
|
||||
from rich.panel import Panel
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
DownloadColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TaskID,
|
||||
TaskProgressColumn,
|
||||
TextColumn,
|
||||
TimeRemainingColumn,
|
||||
TimeElapsedColumn,
|
||||
TransferSpeedColumn,
|
||||
)
|
||||
# rich imports are deferred to avoid ~100ms startup cost.
|
||||
# Classes in this module that use rich types (ProgressBar, PipelineLiveProgress)
|
||||
# import them lazily inside their method bodies at first use.
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -440,6 +428,42 @@ def _sanitise_for_json(
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _import_rich() -> Any:
|
||||
"""Lazy-load rich types used by ProgressBar and PipelineLiveProgress."""
|
||||
import rich.console as _rc
|
||||
import rich.live as _rl
|
||||
import rich.panel as _rp
|
||||
import rich.progress as _rprog
|
||||
# Return a namespace-like object with the types we need
|
||||
class _Rich:
|
||||
Console = _rc.Console
|
||||
ConsoleOptions = _rc.ConsoleOptions
|
||||
Group = _rc.Group
|
||||
Live = _rl.Live
|
||||
Panel = _rp.Panel
|
||||
Progress = _rprog.Progress
|
||||
BarColumn = _rprog.BarColumn
|
||||
DownloadColumn = _rprog.DownloadColumn
|
||||
SpinnerColumn = _rprog.SpinnerColumn
|
||||
TaskID = _rprog.TaskID
|
||||
TaskProgressColumn = _rprog.TaskProgressColumn
|
||||
TextColumn = _rprog.TextColumn
|
||||
TimeRemainingColumn = _rprog.TimeRemainingColumn
|
||||
TimeElapsedColumn = _rprog.TimeElapsedColumn
|
||||
TransferSpeedColumn = _rprog.TransferSpeedColumn
|
||||
return _Rich
|
||||
|
||||
|
||||
_rich: Any = None # cached after first call
|
||||
|
||||
|
||||
def _r() -> Any:
|
||||
global _rich
|
||||
if _rich is None:
|
||||
_rich = _import_rich()
|
||||
return _rich
|
||||
|
||||
|
||||
class ProgressBar:
|
||||
"""Rich progress helper for byte-based transfers.
|
||||
|
||||
@@ -521,16 +545,16 @@ class ProgressBar:
|
||||
console = stderr_console()
|
||||
except Exception:
|
||||
logger.exception("Failed to acquire shared stderr Console from SYS.rich_display; using fallback Console")
|
||||
console = Console(file=stream)
|
||||
console = _r().Console(file=stream)
|
||||
else:
|
||||
console = Console(file=stream)
|
||||
progress = Progress(
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
DownloadColumn(),
|
||||
TransferSpeedColumn(),
|
||||
TimeRemainingColumn(),
|
||||
console = _r().Console(file=stream)
|
||||
progress = _r().Progress(
|
||||
_r().TextColumn("[progress.description]{task.description}"),
|
||||
_r().BarColumn(),
|
||||
_r().TaskProgressColumn(),
|
||||
_r().DownloadColumn(),
|
||||
_r().TransferSpeedColumn(),
|
||||
_r().TimeRemainingColumn(),
|
||||
console=console,
|
||||
transient=True,
|
||||
)
|
||||
@@ -867,7 +891,7 @@ class PipelineLiveProgress:
|
||||
overall = self._overall
|
||||
if pipe_progress is None or transfers is None or overall is None:
|
||||
# Not started (or stopped).
|
||||
yield Panel("", title="Pipeline", expand=False)
|
||||
yield _r().Panel("", title="Pipeline", expand=False)
|
||||
return
|
||||
|
||||
body_parts: List[Any] = [pipe_progress]
|
||||
@@ -875,8 +899,8 @@ class PipelineLiveProgress:
|
||||
body_parts.append(status)
|
||||
body_parts.append(transfers)
|
||||
|
||||
yield Group(
|
||||
Panel(Group(*body_parts),
|
||||
yield _r().Group(
|
||||
_r().Panel(_r().Group(*body_parts),
|
||||
title=self._title_text(),
|
||||
expand=False),
|
||||
overall
|
||||
@@ -895,8 +919,8 @@ class PipelineLiveProgress:
|
||||
if status is not None and self._status_tasks:
|
||||
body_parts.append(status)
|
||||
body_parts.append(transfers)
|
||||
return Group(
|
||||
Panel(Group(*body_parts),
|
||||
return _r().Group(
|
||||
_r().Panel(_r().Group(*body_parts),
|
||||
title=self._title_text(),
|
||||
expand=False),
|
||||
overall
|
||||
@@ -911,55 +935,55 @@ class PipelineLiveProgress:
|
||||
# IMPORTANT: use the shared stderr Console instance so that any
|
||||
# `stderr_console().print(...)` calls from inside cmdlets (e.g. preflight
|
||||
# tables/prompts in download-file) cooperate with Rich Live rendering.
|
||||
# If we create a separate Console(file=sys.stderr), output will fight for
|
||||
# If we create a separate _r().Console(file=sys.stderr), output will fight for
|
||||
# terminal cursor control and appear "blocked"/truncated.
|
||||
from SYS.rich_display import stderr_console
|
||||
|
||||
self._console = stderr_console()
|
||||
|
||||
# Persistent per-pipe bars.
|
||||
self._pipe_progress = Progress(
|
||||
TextColumn("{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
self._pipe_progress = _r().Progress(
|
||||
_r().TextColumn("{task.description}"),
|
||||
_r().TimeElapsedColumn(),
|
||||
_r().BarColumn(),
|
||||
_r().TaskProgressColumn(),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
)
|
||||
|
||||
# Transient, per-item spinner for the currently-active subtask.
|
||||
self._subtasks = Progress(
|
||||
TextColumn(" "),
|
||||
SpinnerColumn("simpleDots"),
|
||||
TextColumn("{task.description}"),
|
||||
self._subtasks = _r().Progress(
|
||||
_r().TextColumn(" "),
|
||||
_r().SpinnerColumn("simpleDots"),
|
||||
_r().TextColumn("{task.description}"),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
)
|
||||
|
||||
# Status line below the pipe bars. Kept simple (no extra bar) so it
|
||||
# doesn't visually offset the main pipe bar columns.
|
||||
self._status = Progress(
|
||||
TextColumn(" [bold]└─ {task.description}[/bold]"),
|
||||
self._status = _r().Progress(
|
||||
_r().TextColumn(" [bold]└─ {task.description}[/bold]"),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
)
|
||||
|
||||
# Byte-based transfer bars (download/upload) integrated into the Live view.
|
||||
self._transfers = Progress(
|
||||
TextColumn(" {task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
DownloadColumn(),
|
||||
TransferSpeedColumn(),
|
||||
TimeRemainingColumn(),
|
||||
self._transfers = _r().Progress(
|
||||
_r().TextColumn(" {task.description}"),
|
||||
_r().BarColumn(),
|
||||
_r().TaskProgressColumn(),
|
||||
_r().DownloadColumn(),
|
||||
_r().TransferSpeedColumn(),
|
||||
_r().TimeRemainingColumn(),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
)
|
||||
|
||||
self._overall = Progress(
|
||||
TimeElapsedColumn(),
|
||||
BarColumn(),
|
||||
TextColumn("{task.description}"),
|
||||
self._overall = _r().Progress(
|
||||
_r().TimeElapsedColumn(),
|
||||
_r().BarColumn(),
|
||||
_r().TextColumn("{task.description}"),
|
||||
console=self._console,
|
||||
transient=False,
|
||||
)
|
||||
@@ -982,7 +1006,7 @@ class PipelineLiveProgress:
|
||||
len(self._pipe_labels)),
|
||||
)
|
||||
|
||||
self._live = Live(
|
||||
self._live = _r().Live(
|
||||
self,
|
||||
console=self._console,
|
||||
refresh_per_second=10,
|
||||
@@ -1011,7 +1035,7 @@ class PipelineLiveProgress:
|
||||
# Not initialized yet; start fresh.
|
||||
self.start()
|
||||
return
|
||||
self._live = Live(
|
||||
self._live = _r().Live(
|
||||
self,
|
||||
console=self._console,
|
||||
refresh_per_second=10,
|
||||
|
||||
Reference in New Issue
Block a user