disks: Add more pretty printing and APIs

Signed-off-by: Ikey Doherty <[email protected]>
AerynOS · Jan 20, 2025 · 078d6ac · 078d6ac
1 parent f7270aa
commit 078d6ac
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 42 deletions.
diff --git a/crates/disks/src/lib.rs b/crates/disks/src/lib.rs
@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: MPL-2.0
 
 use std::{
-    fs, io,
+    fmt, fs, io,
     path::{Path, PathBuf},
 };
 
@@ -54,20 +54,76 @@ pub struct BasicDisk {
     pub partitions: Vec<Partition>,
 }
 
+impl fmt::Display for Disk {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let bytes = self.size();
+        let gib = bytes as f64 / 1_073_741_824.0;
+
+        write!(f, "{} ({:.2} GiB)", self.name(), gib)?;
+
+        if let Some(vendor) = self.vendor() {
+            write!(f, " - {}", vendor)?;
+        }
+
+        if let Some(model) = self.model() {
+            write!(f, " {}", model)?;
+        }
+
+        Ok(())
+    }
+}
+
 impl Disk {
     /// Returns the name of the disk device.
     pub fn name(&self) -> &str {
         match self {
-            Disk::Scsi(disk) => disk.name(),
-            Disk::Nvme(disk) => disk.name(),
+            Disk::Scsi(disk) => &disk.disk.name,
+            Disk::Nvme(disk) => &disk.disk.name,
         }
     }
 
     /// Returns the partitions on the disk.
     pub fn partitions(&self) -> &[Partition] {
         match self {
-            Disk::Scsi(disk) => disk.partitions(),
-            Disk::Nvme(disk) => disk.partitions(),
+            Disk::Scsi(disk) => &disk.disk.partitions,
+            Disk::Nvme(disk) => &disk.disk.partitions,
+        }
+    }
+
+    /// Returns the path to the disk device in dev.
+    pub fn device_path(&self) -> &Path {
+        match self {
+            Disk::Scsi(disk) => &disk.disk.device,
+            Disk::Nvme(disk) => &disk.disk.device,
+        }
+    }
+
+    /// Returns the total number of sectors on the disk.
+    pub fn sectors(&self) -> u64 {
+        match self {
+            Disk::Scsi(disk) => disk.disk.sectors,
+            Disk::Nvme(disk) => disk.disk.sectors,
+        }
+    }
+
+    /// Returns the size of the disk in bytes.
+    pub fn size(&self) -> u64 {
+        self.sectors() * 512
+    }
+
+    /// Returns the model name of the disk.
+    pub fn model(&self) -> Option<&str> {
+        match self {
+            Disk::Scsi(disk) => disk.disk.model.as_deref(),
+            Disk::Nvme(disk) => disk.disk.model.as_deref(),
+        }
+    }
+
+    /// Returns the vendor name of the disk.
+    pub fn vendor(&self) -> Option<&str> {
+        match self {
+            Disk::Scsi(disk) => disk.disk.vendor.as_deref(),
+            Disk::Nvme(disk) => disk.disk.vendor.as_deref(),
         }
     }
 }
@@ -170,9 +226,9 @@ mod tests {
         let devices = BlockDevice::discover().unwrap();
         for device in &devices {
             if let BlockDevice::Disk(disk) = device {
-                println!("{}:", disk.name());
+                println!("{}: {disk}", disk.name());
                 for partition in disk.partitions() {
-                    println!("├─{}", partition.name);
+                    println!("├─{} {partition}", partition.name);
                 }
             }
         }

diff --git a/crates/disks/src/nvme.rs b/crates/disks/src/nvme.rs
@@ -7,7 +7,7 @@
 //! This module provides functionality to enumerate and handle NVMe (Non-Volatile Memory Express)
 //! storage devices by parsing sysfs paths and device names.
 
-use crate::{partition::Partition, BasicDisk, DiskInit};
+use crate::{BasicDisk, DiskInit};
 use regex::Regex;
 use std::{path::Path, sync::OnceLock};
 
@@ -18,7 +18,7 @@ static NVME_PATTERN: OnceLock<Regex> = OnceLock::new();
 #[derive(Debug)]
 pub struct Disk {
     /// The underlying basic disk implementation
-    disk: BasicDisk,
+    pub(crate) disk: BasicDisk,
 }
 
 impl DiskInit for Disk {
@@ -43,18 +43,3 @@ impl DiskInit for Disk {
         }
     }
 }
-
-impl Disk {
-    /// Returns the name of the NVMe disk (e.g. "nvme0n1")
-    ///
-    /// # Returns
-    /// * A string slice containing the disk name
-    pub fn name(&self) -> &str {
-        &self.disk.name
-    }
-
-    /// Returns the partitions on the disk
-    pub fn partitions(&self) -> &[Partition] {
-        &self.disk.partitions
-    }
-}
diff --git a/crates/disks/src/partition.rs b/crates/disks/src/partition.rs
@@ -2,6 +2,7 @@
 //
 // SPDX-License-Identifier: MPL-2.0
 
+use std::fmt;
 use std::path::{Path, PathBuf};
 
 use crate::{sysfs::sysfs_read, DEVFS_DIR, SYSFS_DIR};
@@ -26,6 +27,17 @@ pub struct Partition {
     pub device: PathBuf,
 }
 
+impl fmt::Display for Partition {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "{name} {size:.2} GiB",
+            name = self.name,
+            size = self.size as f64 * 512.0 / (1024.0 * 1024.0 * 1024.0)
+        )
+    }
+}
+
 impl Partition {
     /// Creates a new Partition instance from a sysfs path and partition name.
     ///

diff --git a/crates/disks/src/scsi.rs b/crates/disks/src/scsi.rs
@@ -10,14 +10,14 @@
 
 use std::path::Path;
 
-use crate::{partition::Partition, BasicDisk, DiskInit};
+use crate::{BasicDisk, DiskInit};
 
 /// Represents a SCSI disk device.
 ///
 /// This struct wraps a BasicDisk to provide SCSI-specific functionality.
 #[derive(Debug)]
 pub struct Disk {
-    disk: BasicDisk,
+    pub(crate) disk: BasicDisk,
 }
 
 impl DiskInit for Disk {
@@ -43,19 +43,3 @@ impl DiskInit for Disk {
         }
     }
 }
-
-impl Disk {
-    /// Returns the name of the disk device.
-    ///
-    /// # Returns
-    ///
-    /// The device name (e.g. "sda", "sdb")
-    pub fn name(&self) -> &str {
-        &self.disk.name
-    }
-
-    /// Returns the partitions on the disk.
-    pub fn partitions(&self) -> &[Partition] {
-        &self.disk.partitions
-    }
-}