From 66d3d268444e4f720283a0471f87a8eaf204bded Mon Sep 17 00:00:00 2001
From: Alan Du <alanhdu@gmail.com>
Date: Fri, 30 Aug 2024 15:36:32 -0400
Subject: [PATCH] Fix psutil.cpu_* functions

Use an `@overload` to correctly type the `percpu` argument and
use the right return types
---
 stubs/psutil/psutil/__init__.pyi | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/stubs/psutil/psutil/__init__.pyi b/stubs/psutil/psutil/__init__.pyi
index b01ee35497e0..58d48e3aac92 100644
--- a/stubs/psutil/psutil/__init__.pyi
+++ b/stubs/psutil/psutil/__init__.pyi
@@ -121,12 +121,14 @@ if sys.platform == "win32":
     )
 
 if sys.platform == "linux":
-    from ._pslinux import pfullmem, pmem, sensors_battery as sensors_battery, svmem
+    from ._pslinux import pfullmem, pmem, scputimes, sensors_battery as sensors_battery, svmem
 elif sys.platform == "darwin":
-    from ._psosx import pfullmem, pmem, sensors_battery as sensors_battery, svmem
+    from ._psosx import pfullmem, pmem, scputimes, sensors_battery as sensors_battery, svmem
 elif sys.platform == "win32":
-    from ._pswindows import pfullmem, pmem, sensors_battery as sensors_battery, svmem
+    from ._pswindows import pfullmem, pmem, scputimes, sensors_battery as sensors_battery, svmem
 else:
+    scputimes = Incomplete
+
     class pmem(Any): ...
     class pfullmem(Any): ...
     class svmem(Any): ...
@@ -241,16 +243,27 @@ def wait_procs(
     procs: Iterable[Process], timeout: float | None = None, callback: Callable[[Process], object] | None = None
 ) -> tuple[list[Process], list[Process]]: ...
 def cpu_count(logical: bool = True) -> int: ...
-def cpu_times(percpu: bool = False): ...
+@overload
+def cpu_freq(percpu: Literal[False] = ...) -> scpufreq: ...
+@overload
+def cpu_freq(percpu: Literal[True]) -> list[scpufreq]: ...
+@overload
+def cpu_times(percpu: Literal[False] = ...) -> scputimes: ...
+@overload
+def cpu_times(percpu: Literal[True]) -> list[scputimes]: ...
 @overload
 def cpu_percent(interval: float | None = None, percpu: Literal[False] = False) -> float: ...
 @overload
 def cpu_percent(interval: float | None, percpu: Literal[True]) -> list[float]: ...
 @overload
 def cpu_percent(*, percpu: Literal[True]) -> list[float]: ...
-def cpu_times_percent(interval: float | None = None, percpu: bool = False): ...
+@overload
+def cpu_times_percent(interval: float | None = None, percpu: Literal[False] = False) -> scputimes: ...
+@overload
+def cpu_times_percent(interval: float | None, percpu: Literal[True]) -> list[scputimes]: ...
+@overload
+def cpu_times_percent(*, percpu: Literal[True]) -> list[scputimes]: ...
 def cpu_stats() -> scpustats: ...
-def cpu_freq(percpu: bool = False) -> scpufreq: ...
 def getloadavg() -> tuple[float, float, float]: ...
 def virtual_memory() -> svmem: ...
 def swap_memory() -> sswap: ...