archey4/archey/entries/gpu.py
Michael Bromilow 641dcd21ad
Entries: Use __iter__ and __next__ magic methods.
Incl. changed tests where appropriate.
2024-04-08 22:23:10 +01:00

90 lines
3.1 KiB
Python

"""GPU information detection class"""
import platform
import re
from shlex import split
from subprocess import DEVNULL, CalledProcessError, check_output
from typing import List
from archey.entry import Entry
class GPU(Entry):
"""Relies on `lspci` or `pciconf` to retrieve graphical device(s) information"""
_ICON = "\ue735" # dev_html5_3d_effects
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if platform.system() == "Linux":
self.value = self._parse_lspci_output()
else:
# Darwin or any other BSD-based system.
self.value = self._parse_system_profiler() or self._parse_pciconf_output()
max_count = self.options.get("max_count", 2)
# Consistency with other entries' configuration: Infinite count if false.
if max_count is not False:
self.value = self.value[:max_count]
@staticmethod
def _parse_lspci_output() -> List[str]:
"""Based on `lspci` output, return a list of video controllers names"""
try:
lspci_output = check_output(["lspci", "-m"], universal_newlines=True).splitlines()
except (FileNotFoundError, CalledProcessError):
return []
gpus_list = []
# We'll be looking for specific video controllers (in the below keys order).
for video_key in ("3D", "VGA", "Display"):
for pci_device in lspci_output:
pci_class, pci_vendor, pci_device = split(pci_device)[1:4]
# If a controller type matches the class...
if video_key in pci_class:
# ... adds its name to the final list.
gpus_list.append(f"{pci_vendor} {pci_device}")
return gpus_list
@staticmethod
def _parse_system_profiler() -> List[str]:
"""Based on `system_profiler` output, return a list of video chipsets model names"""
# Parse output from Darwin's `system_profiler` binary.
# We do not use JSON output (more than 10 times longer for nothing).
try:
profiler_output = check_output(
["system_profiler", "SPDisplaysDataType"], stderr=DEVNULL, universal_newlines=True
)
except FileNotFoundError:
return []
return re.findall(r"Chipset Model: (.*)", profiler_output, re.MULTILINE)
@staticmethod
def _parse_pciconf_output() -> List[str]:
"""Based on `pciconf` output, return a list of video devices as long as their vendor"""
try:
pciconf_output = check_output(
["pciconf", "-lv"], stderr=DEVNULL, universal_newlines=True
)
except (FileNotFoundError, CalledProcessError):
return []
gpus_list = []
for vendor, device in re.findall(
r"vendor\s*=\s*\'(.*)\'\s*device\s*=\s*\'(.*)\'\s*class\s*=\s*display",
pciconf_output,
re.MULTILINE,
):
gpus_list.append(f"{vendor} {device}")
return gpus_list
def __next__(self) -> Entry.ValueType:
"""Yield nicely-formatted GPU entry values"""
return (self.name, str(next(self._iter_value)))