Fix mypy on MacOS

This commit is contained in:
Max R 2023-09-09 21:54:47 -04:00
parent 5f4b828999
commit da2b6e6cc4
2 changed files with 19 additions and 20 deletions

View file

@ -24,13 +24,11 @@ TRet = TypeVar('TRet')
def cpu_count() -> int: def cpu_count() -> int:
try:
# On systems that support it, this will return a more accurate count of # On systems that support it, this will return a more accurate count of
# usable CPUs for the current process, which will take into account # usable CPUs for the current process, which will take into account
# cgroup limits # cgroup limits
if hasattr(os, 'sched_getaffinity'):
return len(os.sched_getaffinity(0)) return len(os.sched_getaffinity(0))
except AttributeError:
pass
try: try:
return multiprocessing.cpu_count() return multiprocessing.cpu_count()

View file

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import concurrent.futures import concurrent.futures
import contextlib
import multiprocessing import multiprocessing
import os import os
import sys import sys
@ -20,30 +21,30 @@ def test_cpu_count_sched_getaffinity_exists():
assert xargs.cpu_count() == 345 assert xargs.cpu_count() == 345
@pytest.fixture @contextlib.contextmanager
def no_sched_getaffinity(): def no_sched_getaffinity():
# Simulates an OS without os.sched_getaffinity available (mac/windows) # Simulates an OS without os.sched_getaffinity available (mac/windows)
# https://docs.python.org/3/library/os.html#interface-to-the-scheduler # https://docs.python.org/3/library/os.html#interface-to-the-scheduler
with mock.patch.object( if hasattr(os, 'sched_getaffinity'): # pragma: no cover (win32 py38)
os, func = os.sched_getaffinity
'sched_getaffinity', del os.sched_getaffinity
create=True,
side_effect=AttributeError,
):
yield yield
os.sched_getaffinity = func
else:
yield # pragma: no cover (darwin)
def test_cpu_count_multiprocessing_cpu_count_implemented(no_sched_getaffinity): def test_cpu_count_multiprocessing_cpu_count_implemented():
with mock.patch.object(multiprocessing, 'cpu_count', return_value=123): with mock.patch.object(
multiprocessing, 'cpu_count', return_value=123,
), no_sched_getaffinity():
assert xargs.cpu_count() == 123 assert xargs.cpu_count() == 123
def test_cpu_count_multiprocessing_cpu_count_not_implemented( def test_cpu_count_multiprocessing_cpu_count_not_implemented():
no_sched_getaffinity,
):
with mock.patch.object( with mock.patch.object(
multiprocessing, 'cpu_count', side_effect=NotImplementedError, multiprocessing, 'cpu_count', side_effect=NotImplementedError,
): ), no_sched_getaffinity():
assert xargs.cpu_count() == 1 assert xargs.cpu_count() == 1