Release v0.3.0

Try retrieving physical drive's disk stats before considering partitions

pyudisk/__init__.py

@@ -7,7 +7,7 @@
 
 from .main import EnvConfig, generate_report, monitor, smart_metrics  # noqa: F401
 
-version = "0.2.2"
+version = "0.3.0"
 
 
 @click.command()

pyudisk/config.py

@@ -2,6 +2,7 @@
 import platform
 from typing import Any, List, Optional
 
+from pyarchitecture.config import default_disk_lib
 from pydantic import BaseModel, DirectoryPath, Field, FilePath, HttpUrl, field_validator
 from pydantic_settings import BaseSettings
 
@@ -80,6 +81,11 @@ def get_smart_lib() -> FilePath:
     raise ValueError(f"Unsupported OS: {OPERATING_SYSTEM}")
 
 
+def get_disk_lib() -> str:
+    """Get OS specific disk library to retreive the physical disk ID."""
+    return default_disk_lib()[OPERATING_SYSTEM.lower()]
+
+
 class EnvConfig(BaseSettings):
     """Environment variables configuration.
 
@@ -92,6 +98,7 @@ class EnvConfig(BaseSettings):
     sample_dump: str = "dump.txt"
 
     smart_lib: FilePath = get_smart_lib()
+    disk_lib: FilePath = get_disk_lib()
     metrics: Metric | List[Metric] = Field(default_factory=list)
 
     # Email/SMS notifications

pyudisk/main.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, List, NoReturn
 
 import psutil
+import pyarchitecture.disks
 from psutil._common import sdiskpart
 from pydantic import FilePath, NewPath, ValidationError
 
@@ -18,7 +19,7 @@
 from .util import standard
 
 
-def get_disk(env: EnvConfig) -> Generator[sdiskpart]:
+def get_partitions(env: EnvConfig) -> Generator[sdiskpart]:
     """Gathers disk information using the 'psutil' library.
 
     Args:
@@ -123,7 +124,7 @@ def parse_block_devices(
     category = None
     block_partitions = {
         f"{linux.BlockDevices.head}{block_device.device.split('/')[-1]}:": block_device
-        for block_device in get_disk(env)
+        for block_device in get_partitions(env)
     }
     for line in input_data.splitlines():
         if matching_block := block_partitions.get(line):
@@ -182,20 +183,23 @@ def parse_block_devices(
     return block_devices
 
 
-def get_smart_metrics_macos(smart_lib: FilePath, partition: sdiskpart) -> dict:
+def get_smart_metrics_macos(
+    smart_lib: FilePath, device_id: str, mountpoints: str
+) -> dict:
     """Gathers disk information using the 'smartctl' command on macOS.
 
     Args:
         smart_lib: Library path to 'smartctl' command.
-        partition: Disk partition to gather metrics.
+        device_id: Device ID for the physical drive.
+        mountpoints: Mountpoints for the partitions.
 
     Returns:
         Disk:
         Returns the Disk object with the gathered metrics.
     """
     try:
         result = subprocess.run(
-            [smart_lib, "-a", partition.device, "--json"],
+            [smart_lib, "-a", device_id, "--json"],
             capture_output=True,
             text=True,
             check=False,
@@ -204,7 +208,7 @@ def get_smart_metrics_macos(smart_lib: FilePath, partition: sdiskpart) -> dict:
             LOGGER.debug(
                 "smartctl returned non-zero exit code %d for %s",
                 result.returncode,
-                partition.device,
+                device_id,
             )
         output = json.loads(result.stdout)
     except json.JSONDecodeError as error:
@@ -216,14 +220,21 @@ def get_smart_metrics_macos(smart_lib: FilePath, partition: sdiskpart) -> dict:
         output = {}
     try:
         if output.get("device"):
-            output["usage"] = humanize_usage_metrics(
-                psutil.disk_usage(partition.mountpoint)
-            )
+            if mountpoints:
+                mountpoint = mountpoints[0]
+            else:
+                mountpoint = "/"
+            output["usage"] = humanize_usage_metrics(psutil.disk_usage(mountpoint))
             return output
     except ValidationError as error:
         LOGGER.error(error.errors())
 
 
+def raise_pyarch_error(device: Dict[str, str]) -> NoReturn:
+    """Raises value error for the device specified."""
+    raise ValueError(f"'node' and 'device_id' not found in {device}")
+
+
 def smart_metrics(env: EnvConfig) -> Generator[linux.Disk | darwin.Disk]:
     """Gathers smart metrics using udisksctl dump, and constructs a Disk object.
 
@@ -235,8 +246,27 @@ def smart_metrics(env: EnvConfig) -> Generator[linux.Disk | darwin.Disk]:
         Yields the Disk object from the generated Dataframe.
     """
     if OPERATING_SYSTEM == OperationSystem.darwin:
-        for partition in get_disk(env):
-            if metrics := get_smart_metrics_macos(env.smart_lib, partition):
+        try:
+            disk_data = [
+                (
+                    (
+                        disk.get("node")
+                        or disk.get("device_id")
+                        or raise_pyarch_error(disk)
+                    ),
+                    disk.get("mountpoints"),
+                )
+                for disk in pyarchitecture.disks.get_all_disks(env.disk_lib)
+            ]
+            assert disk_data, "Failed to load physical disk data"
+        except Exception as error:
+            LOGGER.error(error)
+            disk_data = [
+                (partition.device, partition.mountpoint)
+                for partition in get_partitions(env)
+            ]
+        for (device_id, mountpoint) in disk_data:
+            if metrics := get_smart_metrics_macos(env.smart_lib, device_id, mountpoint):
                 try:
                     yield darwin.Disk(**metrics)
                 except ValidationError as error:

requirements.txt

@@ -1,4 +1,5 @@
 click==8.1.*
 psutil==6.0.*
+PyArchitecture>=0.2.0
 pydantic==2.*
 pydantic-settings==2.*