Add progress bar to hassio update/install

.gitignore

@@ -91,6 +91,9 @@ ENV/
 # pylint
 .pylint.d/
 
+#PyCharm
+.idea/
+
 # VS Code
 .vscode/*
 !.vscode/cSpell.json

supervisor/coresys.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypeVar
 
 import aiohttp
@@ -593,7 +594,12 @@ def sys_run_in_executor(
         self, funct: Callable[..., T], *args: Any
     ) -> Coroutine[Any, Any, T]:
         """Add an job to the executor pool."""
-        return self.sys_loop.run_in_executor(None, funct, *args)
+        def funct_with_context():
+            return funct(*args)
+
+        current_context = contextvars.copy_context()
+
+        return self.sys_loop.run_in_executor(None, current_context.run, funct_with_context)
 
     def sys_create_task(self, coroutine: Coroutine) -> asyncio.Task:
         """Create an async task."""

supervisor/docker/__init__.py

@@ -10,6 +10,7 @@
 from packaging import version as pkg_version
 import requests
 
+from .utils import PullProgress
 from ..const import (
     ATTR_REGISTRIES,
     DNS_SUFFIX,
@@ -290,3 +291,18 @@ def check_denylist_images(self) -> bool:
             ", ".join(denied_images),
         )
         return True
+
+    def pull_image(self, image, tag, container_name):
+        """Pull docker image and send progress events to core."""
+        pull = PullProgress(container_name)
+        try:
+            pull.start()
+            pull_log = self.api.pull(image, tag, stream=True, decode=True)
+            for line in pull_log:
+                pull.process_log(line)
+
+            return self.images.get("{0}{2}{1}".format(
+                image, tag, "@" if tag.startswith("sha256:") else ":"
+            ))
+        finally:
+            pull.done()

supervisor/docker/interface.py

@@ -113,7 +113,7 @@ def _install(
             path = IMAGE_WITH_HOST.match(image)
             if path:
                 self._docker_login(path.group(1))
-            docker_image = self.sys_docker.images.pull(f"{image}:{tag}")
+            docker_image = self.sys_docker.pull_image(image, tag, self.name)
             if latest:
                 _LOGGER.info("Tagging image %s with version %s as latest", image, tag)
                 docker_image.tag(image, tag="latest")

supervisor/docker/utils.py

@@ -0,0 +1,97 @@
+"""Utils for Docker."""
+import time
+
+from ..utils import job_monitor
+
+
+class PullProgress:
+    """Docker pull log progress listener."""
+
+    def __init__(self, name: str, sleep=1.0) -> None:
+        """Initialize pull log listener."""
+        self._name = name
+        self._sleep = sleep
+        self._next_push = 0
+        self._downloading = Status()
+        self._extracting = Status()
+        self._job_monitor = job_monitor.get()
+
+    def start(self):
+        """Send progress on start."""
+        self._next_push = time.time()
+        self._send_progress()
+
+    def process_log(self, line):
+        """Process pull log and yield current status."""
+        self._update(line)
+        now = time.time()
+        if self._next_push < now:
+            self._next_push = now + self._sleep
+            self._send_progress()
+
+    def done(self):
+        """Mark current pull as done and send this info to HA Core."""
+        self._downloading.done_all()
+        self._extracting.done_all()
+        self._send_progress()
+
+    def _send_progress(self):
+        if self._job_monitor:
+            self._job_monitor.send_progress(
+                self._name,
+                self._extracting.get(),
+                self._downloading.get(),
+            )
+
+    def _update(self, data):
+        try:
+            layer_id = data["id"]
+            detail = data["progressDetail"]
+            if data["status"] == "Pulling fs layer":
+                # unknown layer size, assume 100MB
+                self._downloading.update(layer_id, 0, 100e6)
+                self._extracting.update(layer_id, 0, 100e6)
+            if data["status"] == "Downloading":
+                self._downloading.update(layer_id, detail["current"], detail["total"])
+                self._extracting.update(layer_id, 0, detail["total"])
+            if data["status"] == "Extracting":
+                self._downloading.done(layer_id)
+                self._extracting.update(layer_id, detail["current"], detail["total"])
+            if data["status"] == "Pull complete":
+                self._downloading.done(layer_id)
+                self._extracting.done(layer_id)
+        except KeyError:
+            pass
+
+
+class Status:
+    """Docker image status object."""
+
+    def __init__(self):
+        """Initialize status object."""
+        self._current = {}
+        self._total = {}
+
+    def update(self, layer_id, current, total):
+        """Update one layer status."""
+        self._current[layer_id] = current
+        self._total[layer_id] = total
+
+    def done(self, layer_id):
+        """Mark one layer as done."""
+        if layer_id in self._total:
+            self._current[layer_id] = self._total[layer_id]
+
+    def done_all(self):
+        """Mark image as done."""
+        if len(self._total) == 0:
+            self.update("id", 1, 1)
+        for layer_id in self._total:
+            self._current[layer_id] = self._total[layer_id]
+
+    def get(self):
+        """Return the current status."""
+        total = sum(self._total.values())
+        if total == 0:
+            return None
+        return sum(self._current.values()) / total

supervisor/utils/__init__.py

@@ -1,16 +1,21 @@
 """Tools file for Supervisor."""
 import asyncio
+from contextvars import ContextVar
 from datetime import datetime
 from ipaddress import IPv4Address
 import logging
 import re
 import socket
 from typing import Any, Optional
 
+from .job_monitor import JobMonitor
+
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 RE_STRING: re.Pattern = re.compile(r"\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))")
 
+job_monitor: ContextVar[Optional[JobMonitor]] = ContextVar("job_monitor", default=None)
+
 
 def convert_to_ascii(raw: bytes) -> str:
     """Convert binary to ascii and remove colors."""
@@ -29,6 +34,8 @@ async def wrap_api(api, *args, **kwargs):
             return False
 
         async with api.lock:
+            job_monitor.set(JobMonitor(api))
+
             return await method(api, *args, **kwargs)
 
     return wrap_api

supervisor/utils/job_monitor.py

@@ -0,0 +1,33 @@
+"""Monitoring class for Supervisor jobs."""
+import asyncio
+from contextlib import suppress
+
+from ..exceptions import HomeAssistantAPIError
+
+
+class JobMonitor:
+    """Monitoring class."""
+
+    def __init__(self, api):
+        self._api = api
+
+    def send_progress(self, name, progress, buffer=None):
+        """Send job progress to core in background task."""
+        self._schedule_send("progress", {
+            "name": name,
+            "progress": progress,
+            "buffer": buffer,
+        })
+
+    def _schedule_send(self, event, json, timeout=2):
+        asyncio.run_coroutine_threadsafe(
+            self._async_send(event, json, timeout),
+            self._api.sys_loop,
+        )
+
+    async def _async_send(self, event, json, timeout) -> None:
+        with suppress(HomeAssistantAPIError):
+            async with self._api.sys_homeassistant.api.make_request(
+                    "post", "api/events/hassio_" + event, json=json, timeout=timeout,
+            ):
+                pass

tests/docker/test_utils.py

@@ -0,0 +1,35 @@
+"""Test Docker Utils."""
+import time
+from unittest.mock import MagicMock, call
+
+from supervisor.docker.utils import PullProgress
+from supervisor.utils import JobMonitor, job_monitor
+from tests.common import load_json_fixture
+
+
+def test_pull_progress():
+    """Test PullProgress class."""
+
+    job = JobMonitor(None)
+    job.send_progress = MagicMock()
+    job_monitor.set(job)
+
+    pull = PullProgress("test-object", 0.01)
+    pull.start()
+    for line in _pull_log_stream():
+        pull.process_log(line)
+    pull.done()
+
+    assert 5 <= len(job.send_progress.mock_calls) <= 7
+
+    first = job.send_progress.mock_calls[0]
+    last = job.send_progress.mock_calls[-1]
+    assert first == call("test-object", None, None)
+    assert last == call("test-object", 1.0, 1.0)
+
+
+def _pull_log_stream():
+    pull_log = load_json_fixture("docker-pull-log.json")
+    for log in pull_log:
+        time.sleep(0.0001)
+        yield log