Improve CpuArch type safety across codebase (#6372)

Co-authored-by: Claude <noreply@anthropic.com>

supervisor/addons/build.py

@@ -20,6 +20,7 @@ from ..const import (
     FILE_SUFFIX_CONFIGURATION,
     META_ADDON,
     SOCKET_DOCKER,
+    CpuArch,
 )
 from ..coresys import CoreSys, CoreSysAttributes
 from ..docker.const import DOCKER_HUB
@@ -67,7 +68,7 @@ class AddonBuild(FileConfiguration, CoreSysAttributes):
         raise RuntimeError()
 
     @cached_property
-    def arch(self) -> str:
+    def arch(self) -> CpuArch:
         """Return arch of the add-on."""
         return self.sys_arch.match([self.addon.arch])
 

supervisor/addons/model.py

@@ -87,6 +87,7 @@ from ..const import (
     AddonBootConfig,
     AddonStage,
     AddonStartup,
+    CpuArch,
 )
 from ..coresys import CoreSys
 from ..docker.const import Capabilities
@@ -548,7 +549,7 @@ class AddonModel(JobGroup, ABC):
         return self.data.get(ATTR_MACHINE, [])
 
     @property
-    def arch(self) -> str:
+    def arch(self) -> CpuArch:
         """Return architecture to use for the addon's image."""
         if ATTR_IMAGE in self.data:
             return self.sys_arch.match(self.data[ATTR_ARCH])

supervisor/arch.py

@@ -4,6 +4,7 @@ import logging
 from pathlib import Path
 import platform
 
+from .const import CpuArch
 from .coresys import CoreSys, CoreSysAttributes
 from .exceptions import ConfigurationFileError, HassioArchNotFound
 from .utils.json import read_json_file
@@ -12,38 +13,40 @@ _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 ARCH_JSON: Path = Path(__file__).parent.joinpath("data/arch.json")
 
-MAP_CPU = {
-    "armv7": "armv7",
-    "armv6": "armhf",
-    "armv8": "aarch64",
-    "aarch64": "aarch64",
-    "i686": "i386",
-    "x86_64": "amd64",
+MAP_CPU: dict[str, CpuArch] = {
+    "armv7": CpuArch.ARMV7,
+    "armv6": CpuArch.ARMHF,
+    "armv8": CpuArch.AARCH64,
+    "aarch64": CpuArch.AARCH64,
+    "i686": CpuArch.I386,
+    "x86_64": CpuArch.AMD64,
 }
 
 
-class CpuArch(CoreSysAttributes):
+class CpuArchManager(CoreSysAttributes):
     """Manage available architectures."""
 
     def __init__(self, coresys: CoreSys) -> None:
         """Initialize CPU Architecture handler."""
         self.coresys = coresys
-        self._supported_arch: list[str] = []
-        self._supported_set: set[str] = set()
-        self._default_arch: str
+        self._supported_arch: list[CpuArch] = []
+        self._supported_set: set[CpuArch] = set()
+        self._default_arch: CpuArch
 
     @property
-    def default(self) -> str:
+    def default(self) -> CpuArch:
         """Return system default arch."""
         return self._default_arch
 
     @property
-    def supervisor(self) -> str:
+    def supervisor(self) -> CpuArch:
         """Return supervisor arch."""
-        return self.sys_supervisor.arch or self._default_arch
+        if self.sys_supervisor.arch:
+            return CpuArch(self.sys_supervisor.arch)
+        return self._default_arch
 
     @property
-    def supported(self) -> list[str]:
+    def supported(self) -> list[CpuArch]:
         """Return support arch by CPU/Machine."""
         return self._supported_arch
 
@@ -65,7 +68,7 @@ class CpuArch(CoreSysAttributes):
             return
 
         # Use configs from arch.json
-        self._supported_arch.extend(arch_data[self.sys_machine])
+        self._supported_arch.extend(CpuArch(a) for a in arch_data[self.sys_machine])
         self._default_arch = self.supported[0]
 
         # Make sure native support is in supported list
@@ -78,14 +81,14 @@ class CpuArch(CoreSysAttributes):
         """Return True if there is a supported arch by this platform."""
         return not self._supported_set.isdisjoint(arch_list)
 
-    def match(self, arch_list: list[str]) -> str:
+    def match(self, arch_list: list[str]) -> CpuArch:
         """Return best match for this CPU/Platform."""
         for self_arch in self.supported:
             if self_arch in arch_list:
                 return self_arch
         raise HassioArchNotFound()
 
-    def detect_cpu(self) -> str:
+    def detect_cpu(self) -> CpuArch:
         """Return the arch type of local CPU."""
         cpu = platform.machine()
         for check, value in MAP_CPU.items():
@@ -96,9 +99,10 @@ class CpuArch(CoreSysAttributes):
                 "Unknown CPU architecture %s, falling back to Supervisor architecture.",
                 cpu,
             )
-            return self.sys_supervisor.arch
+            return CpuArch(self.sys_supervisor.arch)
         _LOGGER.warning(
             "Unknown CPU architecture %s, assuming CPU architecture equals Supervisor architecture.",
             cpu,
         )
-        return cpu
+        # Return the cpu string as-is, wrapped in CpuArch (may fail if invalid)
+        return CpuArch(cpu)

supervisor/bootstrap.py

@@ -13,7 +13,7 @@ from colorlog import ColoredFormatter
 
 from .addons.manager import AddonManager
 from .api import RestAPI
-from .arch import CpuArch
+from .arch import CpuArchManager
 from .auth import Auth
 from .backups.manager import BackupManager
 from .bus import Bus
@@ -71,7 +71,7 @@ async def initialize_coresys() -> CoreSys:
     coresys.jobs = await JobManager(coresys).load_config()
     coresys.core = await Core(coresys).post_init()
     coresys.plugins = await PluginManager(coresys).load_config()
-    coresys.arch = CpuArch(coresys)
+    coresys.arch = CpuArchManager(coresys)
     coresys.auth = await Auth(coresys).load_config()
     coresys.updater = await Updater(coresys).load_config()
     coresys.api = RestAPI(coresys)

supervisor/coresys.py

@@ -29,7 +29,7 @@ from .const import (
 if TYPE_CHECKING:
     from .addons.manager import AddonManager
     from .api import RestAPI
-    from .arch import CpuArch
+    from .arch import CpuArchManager
     from .auth import Auth
     from .backups.manager import BackupManager
     from .bus import Bus
@@ -78,7 +78,7 @@ class CoreSys:
         # Internal objects pointers
         self._docker: DockerAPI | None = None
         self._core: Core | None = None
-        self._arch: CpuArch | None = None
+        self._arch: CpuArchManager | None = None
         self._auth: Auth | None = None
         self._homeassistant: HomeAssistant | None = None
         self._supervisor: Supervisor | None = None
@@ -266,17 +266,17 @@ class CoreSys:
         self._plugins = value
 
     @property
-    def arch(self) -> CpuArch:
-        """Return CpuArch object."""
+    def arch(self) -> CpuArchManager:
+        """Return CpuArchManager object."""
         if self._arch is None:
-            raise RuntimeError("CpuArch not set!")
+            raise RuntimeError("CpuArchManager not set!")
         return self._arch
 
     @arch.setter
-    def arch(self, value: CpuArch) -> None:
-        """Set a CpuArch object."""
+    def arch(self, value: CpuArchManager) -> None:
+        """Set a CpuArchManager object."""
         if self._arch:
-            raise RuntimeError("CpuArch already set!")
+            raise RuntimeError("CpuArchManager already set!")
         self._arch = value
 
     @property
@@ -733,8 +733,8 @@ class CoreSysAttributes:
         return self.coresys.plugins
 
     @property
-    def sys_arch(self) -> CpuArch:
-        """Return CpuArch object."""
+    def sys_arch(self) -> CpuArchManager:
+        """Return CpuArchManager object."""
         return self.coresys.arch
 
     @property

supervisor/docker/interface.py

@@ -52,7 +52,7 @@ from .stats import DockerStats
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
-MAP_ARCH: dict[CpuArch | str, str] = {
+MAP_ARCH: dict[CpuArch, str] = {
     CpuArch.ARMV7: "linux/arm/v7",
     CpuArch.ARMHF: "linux/arm/v6",
     CpuArch.AARCH64: "linux/arm64",
@@ -366,7 +366,7 @@ class DockerInterface(JobGroup, ABC):
         if not image:
             raise ValueError("Cannot pull without an image!")
 
-        image_arch = str(arch) if arch else self.sys_arch.supervisor
+        image_arch = arch or self.sys_arch.supervisor
         listener: EventListener | None = None
 
         _LOGGER.info("Downloading docker image %s with tag %s.", image, version)
@@ -603,9 +603,7 @@ class DockerInterface(JobGroup, ABC):
         expected_cpu_arch: CpuArch | None = None,
     ) -> None:
         """Check we have expected image with correct arch."""
-        expected_image_cpu_arch = (
-            str(expected_cpu_arch) if expected_cpu_arch else self.sys_arch.supervisor
-        )
+        arch = expected_cpu_arch or self.sys_arch.supervisor
         image_name = f"{expected_image}:{version!s}"
         if self.image == expected_image:
             try:
@@ -623,7 +621,7 @@ class DockerInterface(JobGroup, ABC):
             # If we have an image and its the right arch, all set
             # It seems that newer Docker version return a variant for arm64 images.
             # Make sure we match linux/arm64 and linux/arm64/v8.
-            expected_image_arch = MAP_ARCH[expected_image_cpu_arch]
+            expected_image_arch = MAP_ARCH[arch]
             if image_arch.startswith(expected_image_arch):
                 return
             _LOGGER.info(
@@ -636,7 +634,7 @@ class DockerInterface(JobGroup, ABC):
         # We're missing the image we need. Stop and clean up what we have then pull the right one
         with suppress(DockerError):
             await self.remove()
-        await self.install(version, expected_image, arch=expected_image_cpu_arch)
+        await self.install(version, expected_image, arch=arch)
 
     @Job(
         name="docker_interface_update",

tests/addons/test_manager.py

@@ -10,7 +10,7 @@ from awesomeversion import AwesomeVersion
 import pytest
 
 from supervisor.addons.addon import Addon
-from supervisor.arch import CpuArch
+from supervisor.arch import CpuArchManager
 from supervisor.config import CoreConfig
 from supervisor.const import AddonBoot, AddonStartup, AddonState, BusEvent
 from supervisor.coresys import CoreSys
@@ -54,7 +54,9 @@ async def fixture_mock_arch_disk() -> AsyncGenerator[None]:
     """Mock supported arch and disk space."""
     with (
         patch("shutil.disk_usage", return_value=(42, 42, 2 * (1024.0**3))),
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
     ):
         yield
 

tests/api/test_addons.py

@@ -9,7 +9,7 @@ import pytest
 
 from supervisor.addons.addon import Addon
 from supervisor.addons.build import AddonBuild
-from supervisor.arch import CpuArch
+from supervisor.arch import CpuArchManager
 from supervisor.const import AddonState
 from supervisor.coresys import CoreSys
 from supervisor.docker.addon import DockerAddon
@@ -236,7 +236,9 @@ async def test_api_addon_rebuild_healthcheck(
         patch.object(AddonBuild, "is_valid", return_value=True),
         patch.object(DockerAddon, "is_running", return_value=False),
         patch.object(Addon, "need_build", new=PropertyMock(return_value=True)),
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
         patch.object(DockerAddon, "run", new=container_events_task),
         patch.object(
             coresys.docker,
@@ -308,7 +310,9 @@ async def test_api_addon_rebuild_force(
         patch.object(
             Addon, "need_build", new=PropertyMock(return_value=False)
         ),  # Image-based
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
     ):
         resp = await api_client.post("/addons/local_ssh/rebuild")
 
@@ -326,7 +330,9 @@ async def test_api_addon_rebuild_force(
         patch.object(
             Addon, "need_build", new=PropertyMock(return_value=False)
         ),  # Image-based
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
         patch.object(DockerAddon, "run", new=container_events_task),
         patch.object(
             coresys.docker,

tests/api/test_store.py

@@ -10,7 +10,7 @@ from awesomeversion import AwesomeVersion
 import pytest
 
 from supervisor.addons.addon import Addon
-from supervisor.arch import CpuArch
+from supervisor.arch import CpuArchManager
 from supervisor.backups.manager import BackupManager
 from supervisor.config import CoreConfig
 from supervisor.const import AddonState, CoreState
@@ -191,7 +191,9 @@ async def test_api_store_update_healthcheck(
         patch.object(DockerAddon, "run", new=container_events_task),
         patch.object(DockerInterface, "install"),
         patch.object(DockerAddon, "is_running", return_value=False),
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
     ):
         resp = await api_client.post(f"/store/addons/{TEST_ADDON_SLUG}/update")
 
@@ -548,7 +550,9 @@ async def test_api_store_addons_addon_availability_arch_not_supported(
         coresys.addons.data.user[addon_obj.slug] = {"version": AwesomeVersion("0.0.1")}
 
     # Mock the system architecture to be different
-    with patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])):
+    with patch.object(
+        CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+    ):
         resp = await api_client.request(
             api_method, f"/store/addons/{addon_obj.slug}/{api_action}"
         )

tests/store/test_store_manager.py

@@ -9,7 +9,7 @@ from awesomeversion import AwesomeVersion
 import pytest
 
 from supervisor.addons.addon import Addon
-from supervisor.arch import CpuArch
+from supervisor.arch import CpuArchManager
 from supervisor.backups.manager import BackupManager
 from supervisor.coresys import CoreSys
 from supervisor.exceptions import AddonNotSupportedError, StoreJobError
@@ -163,7 +163,9 @@ async def test_update_unavailable_addon(
     with (
         patch.object(BackupManager, "do_backup_partial") as backup,
         patch.object(AddonStore, "data", new=PropertyMock(return_value=addon_config)),
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
         patch.object(CoreSys, "machine", new=PropertyMock(return_value="qemux86-64")),
         patch.object(
             HomeAssistant,
@@ -219,7 +221,9 @@ async def test_install_unavailable_addon(
 
     with (
         patch.object(AddonStore, "data", new=PropertyMock(return_value=addon_config)),
-        patch.object(CpuArch, "supported", new=PropertyMock(return_value=["amd64"])),
+        patch.object(
+            CpuArchManager, "supported", new=PropertyMock(return_value=["amd64"])
+        ),
         patch.object(CoreSys, "machine", new=PropertyMock(return_value="qemux86-64")),
         patch.object(
             HomeAssistant,