From 605b40a8f71ba122503ce455a974edeebcf7abd7 Mon Sep 17 00:00:00 2001
From: Filip Strajnar <filip.strajnar@gmail.com>
Date: Sat, 26 Oct 2024 20:40:20 +0200
Subject: [PATCH] Added methods vendor_code and device_code.

---
 src/pci_passthrough_assist/pci_device.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/pci_passthrough_assist/pci_device.py b/src/pci_passthrough_assist/pci_device.py
index 491d81d..e5539a2 100644
--- a/src/pci_passthrough_assist/pci_device.py
+++ b/src/pci_passthrough_assist/pci_device.py
@@ -14,6 +14,16 @@ class PciDevice:
     def is_vga(self) -> bool:
         return exists(f"/sys/bus/pci/devices/{self.device_id}/boot_vga")
 
+    def vendor_code(self, remove_prefix: bool = True) -> str:
+        with open(f"/sys/bus/pci/devices/{self.device_id}/vendor") as vendor:
+            return vendor.read() if not remove_prefix else vendor.read(
+            ).lstrip("0x")
+
+    def device_code(self, remove_prefix: bool = True) -> str:
+        with open(f"/sys/bus/pci/devices/{self.device_id}/vendor") as device:
+            return device.read() if not remove_prefix else device.read(
+            ).lstrip("0x")
+
     def unbind_driver(self):
         with open(f"/sys/bus/pci/devices/{self.device_id}/driver/unbind",
                   "w") as device_driver: