feat: feat(print): copy USB gcode to local storage and print the local copy

BlocksScreen/lib/panels/printTab.py

@@ -1,8 +1,10 @@
 import logging
 import os
+import pathlib
 import typing
 from functools import partial
 
+import helper_methods
 from configfile import BlocksScreenConfig, get_configparser
 from lib.files import Files
 from lib.moonrakerComm import MoonWebSocket
@@ -20,11 +22,18 @@
 from lib.utils import gcode_loader
 from lib.utils.blocks_button import BlocksCustomButton
 from lib.utils.display_button import DisplayButton
+from lib.utils.usb_import import UsbImportTask, remember_import
 from PyQt6 import QtCore, QtGui, QtWidgets
 
 if typing.TYPE_CHECKING:
     from devices.amu import AMUManager
 
+USB_SETTLE_TIMEOUT_MS = 120_000  # Happy Hare pre-processing of a large file
+# RF50-Klipper's SDCARD_PRINT_FILE override emits this for USB files (Mainsail starts).
+USB_PRINT_ACTION = "// action:blocks_usb_print "
+# The override hands USB files over only once this is True; older pairs print directly.
+USB_PRINT_FLAG = "blocksscreen_usb"
+
 logger = logging.getLogger(__name__)
 
 
@@ -78,6 +87,15 @@ def __init__(
         self._finish_print_handled: bool = False
         self._cancel_z_snapshot: float = 0.0
         self._z_apply_command: str = "Z_OFFSET_APPLY_ENDSTOP"
+        self._usb_import_task: UsbImportTask | None = None
+        self._usb_key: str = ""
+        self._usb_local: str = ""
+        self._usb_created: set[str] = set()
+        # Moonraker announces a new gcode only once its pre-processing is done.
+        self._usb_settle_timer = QtCore.QTimer(self)
+        self._usb_settle_timer.setSingleShot(True)
+        self._usb_settle_timer.setInterval(USB_SETTLE_TIMEOUT_MS)
+        self._usb_settle_timer.timeout.connect(self._on_usb_settle_timeout)
 
         self.setupMainPrintPage()
         self.ws: MoonWebSocket = ws
@@ -138,6 +156,7 @@ def __init__(
         self.filesPage_widget.request_dir_info.connect(self.file_data.request_dir_info)
         self.file_data.on_file_list.connect(self.filesPage_widget.on_file_list)
         self.file_data.file_added.connect(self.filesPage_widget.on_file_added)
+        self.file_data.file_added.connect(self._on_usb_file_added)
         self.file_data.file_removed.connect(self.filesPage_widget.on_file_removed)
         self.file_data.file_modified.connect(self.filesPage_widget.on_file_modified)
         self.file_data.dir_added.connect(self.filesPage_widget.on_dir_added)
@@ -209,6 +228,8 @@ def __init__(
             self.jobStatusPage_widget.on_gcode_move_update
         )
         self.printer.request_available_objects_signal.connect(self.klipper_ready_signal)
+        self.printer.gcode_response.connect(self._on_gcode_response)
+        self.printer.gcode_macro_update.connect(self._on_gcode_macro_update)
         self.babystepPage = BabystepPage(self)
         self.babystepPage.request_back.connect(self.back_button)
         self.addWidget(self.babystepPage)
@@ -359,8 +380,95 @@ def start_print(self, file: str) -> None:
             logger.warning("Print start of %s blocked: %s", file, reason)
             self.notify.emit("PrintTab", reason, 2, True)
             return
+        if helper_methods.is_usb_path(file):
+            self._start_usb_import(file)
+            return
         self.jobStatusPage_widget.on_print_start(file)
 
+    @QtCore.pyqtSlot(list, name="on_gcode_response")
+    def _on_gcode_response(self, messages: list) -> None:
+        """Import a USB file whose print another client started through Klipper."""
+        for line in messages:
+            file = str(line).removeprefix(USB_PRINT_ACTION)
+            # The macro sends USB paths only; a local one would bypass the touchscreen.
+            if file != line and helper_methods.is_usb_path(file):
+                # Ack first: an unacked hand-off makes the macro print from USB itself.
+                self.run_gcode_signal.emit(
+                    "SET_GCODE_VARIABLE MACRO=SDCARD_PRINT_FILE"
+                    " VARIABLE=usb_pending VALUE='\"\"'"
+                )
+                self.start_print(file)
+
+    @QtCore.pyqtSlot(str, dict, name="on_gcode_macro_update")
+    def _on_gcode_macro_update(self, name: str, values: dict) -> None:
+        """Tell the RF50-Klipper USB override that BlocksScreen handles its files."""
+        # False only on the override that knows the flag; Klipper resets it on restart.
+        if name.upper() == "SDCARD_PRINT_FILE" and values.get(USB_PRINT_FLAG) is False:
+            self.run_gcode_signal.emit(
+                f"SET_GCODE_VARIABLE MACRO=SDCARD_PRINT_FILE VARIABLE={USB_PRINT_FLAG}"
+                " VALUE=True"
+            )
+
+    def _start_usb_import(self, file: str) -> None:
+        """Copy a USB gcode to the gcodes root off-thread, then print the copy."""
+        # Printing from the stick would fail mid-job if it is pulled or glitches.
+        if self._usb_import_task is not None:
+            logger.info("USB import already running, ignoring %s", file)
+            return
+        self._usb_created.clear()
+        self._usb_import_task = UsbImportTask(pathlib.Path(self.gcode_path), file)
+        self._usb_import_task.signals.done.connect(self._on_usb_import_done)
+        self.call_load_panel.emit(True, "Copying file from USB...\nPlease wait", False)
+        QtCore.QThreadPool.globalInstance().start(self._usb_import_task)
+
+    @QtCore.pyqtSlot(str, str, str, bool, str, name="on_usb_import_done")
+    def _on_usb_import_done(
+        self, usb_file: str, local_file: str, key: str, reused: bool, error: str
+    ) -> None:
+        """Print a reused copy now, a new one once Moonraker has processed it."""
+        if error:
+            self._usb_import_task = None
+            self.call_load_panel.emit(False, "", False)
+            self.notify.emit("PrintTab", f"USB print failed: {error}", 2, True)
+            return
+        logger.info("Local copy of %s is %s", usb_file, local_file)
+        self._usb_key, self._usb_local = key, local_file
+        # Printing before Happy Hare rewrites the file runs the unprocessed gcode.
+        if reused or local_file in self._usb_created:
+            self._usb_print_local_copy()
+        else:
+            self._usb_settle_timer.start()
+
+    @QtCore.pyqtSlot(dict, name="on_usb_file_added")
+    def _on_usb_file_added(self, item: dict) -> None:
+        """Release a waiting USB print when Moonraker announces its local copy."""
+        if self._usb_import_task is None:
+            return
+        name = str(item.get("path", "")).removeprefix("/")
+        self._usb_created.add(name)  # may beat the worker's done signal
+        if name == self._usb_local:
+            self._usb_print_local_copy()
+
+    @QtCore.pyqtSlot(name="on_usb_settle_timeout")
+    def _on_usb_settle_timeout(self) -> None:
+        """Print anyway if Moonraker never announced the copy."""
+        logger.warning("No create_file for %s, printing it as is", self._usb_local)
+        self._usb_print_local_copy()
+
+    def _usb_print_local_copy(self) -> None:
+        """Index the settled local copy and start printing it."""
+        self._usb_settle_timer.stop()
+        local = self._usb_local
+        try:
+            remember_import(pathlib.Path(self.gcode_path), self._usb_key, local)
+        except OSError as e:
+            logger.warning("Could not index USB copy %s: %s", local, e)
+        self._usb_import_task = None
+        self._usb_key = self._usb_local = ""
+        self.call_load_panel.emit(False, "", False)
+        logger.info("Printing local copy %s", local)
+        self.jobStatusPage_widget.on_print_start(local)
+
     @QtCore.pyqtSlot(str, str, name="delete_file")
     @QtCore.pyqtSlot(str, name="delete_file")
     def delete_file(self, filename: str, directory: str = "gcodes") -> None:

BlocksScreen/lib/utils/usb_import.py

@@ -0,0 +1,214 @@
+"""Copies USB gcode files into local storage so prints survive drive removal."""
+
+import filecmp
+import hashlib
+import json
+import logging
+import os
+import pathlib
+import shutil
+
+import helper_methods
+from PyQt6 import QtCore
+
+logger = logging.getLogger(__name__)
+
+MB = 1024 * 1024
+ALLOWED_EXT = ".gcode"  # the only extension lib.files lists and notifies
+MIN_FREE_MARGIN = 64 * MB  # free space left on the SD card after the copy
+SAMPLE = 64 * 1024
+STATE_DIR_NAME = ".usb_import"
+INDEX_NAME = "index.json"
+
+
+class UsbImportError(Exception):
+    """Import failure whose message is safe to show on screen."""
+
+
+def _state_dir(gcode_root: pathlib.Path) -> pathlib.Path:
+    """Index and scratch dir, beside the gcodes tree so Moonraker never lists it."""
+    return gcode_root.parent / STATE_DIR_NAME
+
+
+def _source_key(src: pathlib.Path, size: int) -> str:
+    """Identity of a USB file from its name, size, head and tail."""
+    # Slicers stamp the header and append their config, so ends differ per slice.
+    h = hashlib.sha256(f"{src.name}\0{size}\0".encode())
+    with src.open("rb") as f:
+        h.update(f.read(SAMPLE))
+        if size > SAMPLE:
+            f.seek(max(SAMPLE, size - SAMPLE))
+            h.update(f.read(SAMPLE))
+    return h.hexdigest()
+
+
+def _load_index(gcode_root: pathlib.Path) -> dict:
+    """Source key -> {file, size, mtime_ns} of earlier imports; empty if unreadable."""
+    try:
+        data = json.loads((_state_dir(gcode_root) / INDEX_NAME).read_text())
+    except (OSError, ValueError):
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def _indexed_copy(gcode_root: pathlib.Path, key: str) -> pathlib.Path | None:
+    """The earlier local copy for *key*, if it is still there and untouched."""
+    entry = _load_index(gcode_root).get(key)
+    if not isinstance(entry, dict) or not isinstance(entry.get("file"), str):
+        return None
+    local = gcode_root / entry["file"]
+    try:
+        st = local.stat()
+    except OSError:
+        return None
+    if (st.st_size, st.st_mtime_ns) != (entry.get("size"), entry.get("mtime_ns")):
+        return None  # replaced or edited since, so it may no longer be this file
+    return local
+
+
+def remember_import(gcode_root: pathlib.Path, key: str, local_name: str) -> None:
+    """Record *local_name* as the copy for *key*, once Moonraker has processed it."""
+    # Moonraker and Happy Hare rewrite new gcodes, so the copy is only final now.
+    root = gcode_root.resolve()
+    st = (root / local_name).stat()
+    index = {
+        k: v
+        for k, v in _load_index(root).items()
+        if isinstance(v, dict) and (root / str(v.get("file"))).is_file()
+    }
+    index[key] = {"file": local_name, "size": st.st_size, "mtime_ns": st.st_mtime_ns}
+    state = _state_dir(root)
+    state.mkdir(exist_ok=True)
+    tmp = state / f".{INDEX_NAME}.part"
+    tmp.write_text(json.dumps(index, indent=1))
+    os.replace(tmp, state / INDEX_NAME)
+
+
+def _tmp_dir(gcode_root: pathlib.Path) -> pathlib.Path:
+    """Scratch dir outside the gcodes tree, or the gcodes root if on another fs."""
+    # Outside Moonraker's inotify tree the final rename reads as a plain create_file.
+    tmp_dir = _state_dir(gcode_root)
+    try:
+        tmp_dir.mkdir(exist_ok=True)
+        if os.stat(tmp_dir).st_dev == os.stat(gcode_root).st_dev:
+            return tmp_dir  # same fs keeps os.replace atomic
+    except OSError:
+        pass
+    return gcode_root
+
+
+def _pick_destination(
+    src: pathlib.Path, size: int, dest_dir: pathlib.Path
+) -> tuple[pathlib.Path, bool]:
+    """(path, reused): a byte-identical existing file, else the first free name."""
+    n = 0
+    while True:
+        cand = dest_dir / (src.name if n == 0 else f"{src.stem}_{n}{src.suffix}")
+        if not cand.exists() and not cand.is_symlink():
+            return cand, False
+        if (
+            cand.is_file()
+            and cand.stat().st_size == size
+            and filecmp.cmp(src, cand, shallow=False)
+        ):
+            return cand, True
+        n += 1  # same name, other content: never overwrite a local file
+
+
+def _fsync(path: pathlib.Path) -> None:
+    """Flush *path* (file or directory) to disk."""
+    fd = os.open(path, os.O_RDONLY)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _copy_atomic(
+    src: pathlib.Path,
+    src_st: os.stat_result,
+    final: pathlib.Path,
+    tmp_dir: pathlib.Path,
+) -> None:
+    """Copy *src* to *final* via a synced temp file, so *final* is never partial."""
+    for stale in tmp_dir.glob(".*.part"):  # leftovers from a crash or power loss
+        stale.unlink(missing_ok=True)
+    free = shutil.disk_usage(final.parent).free
+    if free < src_st.st_size + MIN_FREE_MARGIN:
+        raise UsbImportError(f"Not enough free space ({free // MB} MB free)")
+    tmp = tmp_dir / f".{final.name}.part"
+    try:
+        shutil.copyfile(src, tmp)  # kernel-side copy on Linux
+        now = src.stat()
+        if (now.st_size, now.st_mtime_ns) != (src_st.st_size, src_st.st_mtime_ns):
+            raise UsbImportError("The file changed during the copy")
+        if tmp.stat().st_size != src_st.st_size:
+            raise UsbImportError("Copy incomplete")
+        _fsync(tmp)
+        os.replace(tmp, final)
+        _fsync(final.parent)  # persist the rename itself across power loss
+    finally:
+        tmp.unlink(missing_ok=True)
+
+
+def import_usb_gcode(gcode_root: pathlib.Path, filename: str) -> tuple[str, bool, str]:
+    """Copy USB *filename* into *gcode_root*; returns (local name, reused, key)."""
+    root = gcode_root.resolve()
+    rel = filename.strip("/")
+    if not helper_methods.is_usb_path(rel):
+        raise UsbImportError("Not a USB file")
+    usb_root = (root / rel.split("/", 1)[0]).resolve()
+    src = (root / rel).resolve()
+    if usb_root not in src.parents:  # symlinks or ".." escaping the drive
+        raise UsbImportError("Invalid file path")
+    if src.suffix.lower() != ALLOWED_EXT:
+        raise UsbImportError("Unsupported file type")
+    if not os.path.ismount(usb_root) or not src.is_file():
+        raise UsbImportError("File not found, was the USB drive removed?")
+    try:
+        src_st = src.stat()
+        key = _source_key(src, src_st.st_size)
+        final, reused = _indexed_copy(root, key), True
+        if final is None:
+            final, reused = _pick_destination(src, src_st.st_size, root)
+        if not reused:
+            _copy_atomic(src, src_st, final, _tmp_dir(root))
+    except OSError as e:
+        raise UsbImportError(
+            f"Copy failed, was the USB drive removed? ({e.strerror})"
+        ) from e
+    logger.info("USB import %s -> %s (reused=%s)", rel, final.name, reused)
+    return final.name, reused, key
+
+
+class _ImportSignals(QtCore.QObject):
+    """Worker result signal; QRunnable can't hold signals."""
+
+    # (usb_file, local_file, key, reused, error)
+    done = QtCore.pyqtSignal(str, str, str, bool, str)
+
+
+class UsbImportTask(QtCore.QRunnable):
+    """Runs import_usb_gcode off the UI thread."""
+
+    def __init__(self, gcode_root: pathlib.Path, filename: str) -> None:
+        super().__init__()
+        self.setAutoDelete(False)  # the caller holds it until done fires
+        self._gcode_root = gcode_root
+        self._filename = filename
+        self.signals = _ImportSignals()
+
+    @QtCore.pyqtSlot()
+    def run(self) -> None:
+        """Copy, then report the local copy or the error text."""
+        try:
+            local, reused, key = import_usb_gcode(self._gcode_root, self._filename)
+            self.signals.done.emit(self._filename, local, key, reused, "")
+        except UsbImportError as e:
+            logger.warning("USB import of %s failed: %s", self._filename, e)
+            self.signals.done.emit(self._filename, "", "", False, str(e))
+        except Exception:  # never leave the loadscreen up on a bug
+            logger.exception("USB import of %s crashed", self._filename)
+            self.signals.done.emit(
+                self._filename, "", "", False, "Unexpected error copying file"
+            )