# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['CPU_COUNT', 'SCHED_GETAFFINITY', 'mod_cpu_count', 'getaffinity0', 'os_cpu_count', 'jl_cpu_count', 'mp_cpu_count',
           'numcores']

# %% ../nbs/00_core.ipynb 6
# import os, joblib, multiprocessing
from importlib import import_module
from inspect import signature

# %% ../nbs/00_core.ipynb 8
from types import ModuleType

# %% ../nbs/00_core.ipynb 11
#| export

# %% ../nbs/00_core.ipynb 13
#| export


# %% ../nbs/00_core.ipynb 15
#| export


# %% ../nbs/00_core.ipynb 17
#| export


# %% ../nbs/00_core.ipynb 19
#| export


# %% ../nbs/00_core.ipynb 22
CPU_COUNT = 'cpu_count'
SCHED_GETAFFINITY = 'sched_getaffinity'

def mod_cpu_count(module: str | ModuleType, method: str = CPU_COUNT, *args, **kwargs) -> int:
    if isinstance(module, str): module = import_module(module)
    func = getattr(module, method)
    fsig = signature(func)
    barg = fsig.bind(*args, **kwargs)
    return func(*barg.args, **barg.kwargs)

def getaffinity0() -> int:
    import os
    if hasattr(os, SCHED_GETAFFINITY):
        cpus = os.sched_getaffinity(0)
        return len(cpus)
    return 0

def os_cpu_count() -> int:
    return getaffinity0() or mod_cpu_count('os')

def jl_cpu_count() -> int:
    return mod_cpu_count('joblib')

def mp_cpu_count():
    return mod_cpu_count('multiprocessing')

# %% ../nbs/00_core.ipynb 23
def numcores(module: str | ModuleType = 'os'):
    if isinstance(module, ModuleType): module = module.__name__
    match module:
        case 'os': return os_cpu_count()
        case 'jl': return jl_cpu_count()
        case 'mp': return mp_cpu_count()
        case 'joblib': return jl_cpu_count()
        case 'multiprocessing': return mp_cpu_count()
        case _: return os_cpu_count()
    return 0

# %% ../nbs/00_core.ipynb 25
#| export
