[python-package] fix mypy error about cpu_count() methods (#5786)

a6bf4ad4 · James Lamb · GitHub · 54486b4f · a6bf4ad4
Unverified Commit a6bf4ad4 authored Mar 16, 2023 by James Lamb Committed by GitHub Mar 16, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

python-package/lightgbm/compat.py python-package/lightgbm/compat.py +4 -4

No files found.
--- a/python-package/lightgbm/compat.py
+++ b/python-package/lightgbm/compat.py
@@ -183,18 +183,18 @@ except ImportError:
 try:
    from joblib import cpu_count

-    def _LGBMCpuCount(only_physical_cores: bool = True):
+    def _LGBMCpuCount(only_physical_cores: bool = True) -> int:
        return cpu_count(only_physical_cores=only_physical_cores)
 except ImportError:
    try:
        from psutil import cpu_count

-        def _LGBMCpuCount(only_physical_cores: bool = True):
-            return cpu_count(logical=not only_physical_cores)
+        def _LGBMCpuCount(only_physical_cores: bool = True) -> int:
+            return cpu_count(logical=not only_physical_cores) or 1
    except ImportError:
        from multiprocessing import cpu_count

-        def _LGBMCpuCount(only_physical_cores: bool = True):
+        def _LGBMCpuCount(only_physical_cores: bool = True) -> int:
            return cpu_count()

 __all__: List[str] = []