Inspect/config numpy-linked OpenBLAS

numpy_openblas_config.py

"""
Inspect/config numpy-linked OpenBLAS
"""

import re
import ctypes

import numpy.core


class openblas_config:
    def get_config(self) -> bytes: ...
    def get_num_threads(self) -> int: ...
    def set_num_threads(self, i:int): ...


def numpy_openblas():
    # type: () -> Optional[openblas_config]
    try: 
        handle = ctypes.CDLL(numpy.core.multiarray.__file__) 
    except OSError: 
        handle = ctypes.CDLL(numpy.core._multiarray_umath.__file__)  
    try:
        openblas_get_config = handle["openblas_get_config"]
        openblas_get_num_threads = handle["openblas_get_num_threads"]
        openblas_set_num_threads = handle["openblas_set_num_threads"]
    except AttributeError:
        return None
    openblas_get_config.argtypes = []
    openblas_get_config.restype = ctypes.c_char_p  # pointer to static memory
    openblas_get_num_threads.argtypes = []
    openblas_get_num_threads.restype = ctypes.c_int
    openblas_set_num_threads.argtypes = [ctypes.c_int]
    openblas_set_num_threads.restype = None

    config = openblas_config()
    config.get_config = openblas_get_config
    config.set_num_threads = openblas_set_num_threads
    config.get_num_threads = openblas_get_num_threads
    return config


def numpy_openblas_get_config_dict():
    # type: () -> Optional[Dict[str, Any]]
    cfg = numpy_openblas()
    if cfg is None:
        return None
    config = cfg.get_config()
    try:
        match = next(re.finditer(br"MAX_THREADS=(\d+)", config))
    except StopIteration:
        return {"MAX_THREADS": 1}
    else:
        return {"MAX_THREADS": int(match.group(1))}