setup.py

#!/usr/bin/env python3
###################################################################
#  NumExpr - Fast numerical array expression evaluator for NumPy.
#
#      License: MIT
#      Author:  See AUTHORS.txt
#
#  See LICENSE.txt and LICENSES/*.txt for details about copyright and
#  rights to use.
####################################################################

import os, os.path as op
from setuptools import setup, Extension
from setuptools.config import read_configuration
import platform
import numpy as np

with open('requirements.txt') as f:
    requirements = f.read().splitlines()

with open('numexpr/version.py', 'w') as fh:
    cfg = read_configuration("setup.cfg")
    fh.write('# THIS FILE IS GENERATED BY `SETUP.PY`\n')
    fh.write("version = '%s'\n" % cfg['metadata']['version'])
    try:
        fh.write("numpy_build_version = '%s'\n" % np.__version__)
    except ImportError:
        pass
    fh.write("platform_machine = '%s'\n" % platform.machine())

lib_dirs = []
inc_dirs = [np.get_include()]
libs = []  # Pre-built libraries ONLY, like python36.so
clibs = []
def_macros = []
sources = ['numexpr/interpreter.cpp',
           'numexpr/module.cpp',
           'numexpr/numexpr_object.cpp']
extra_cflags = []
extra_link_args = []

if platform.uname().system == 'Windows':
    # For MSVC only
    if "MSC" in platform.python_compiler():
        extra_cflags = ['/O2']
    extra_link_args = []
    sources.append('numexpr/win32/pthread.c')
else:
    extra_cflags = []
    extra_link_args = []


def parse_site_cfg():
    """
    Parses `site.cfg`, if it exists, to determine the location of Intel oneAPI MKL.

    To compile NumExpr with MKL (VML) support, typically you need to copy the 
    provided `site.cfg.example` to `site.cfg` and then edit the paths in the 
    configuration lines for `include_dirs` and `library_dirs` paths to point 
    to the appropriate directories on your machine.
    """
    import configparser
    site = configparser.ConfigParser()
    if not op.isfile('site.cfg'):
        return

    site.read('site.cfg')

    if 'mkl' in site.sections():
        inc_dirs.extend(
            site['mkl']['include_dirs'].split(os.pathsep))
        lib_dirs.extend(
            site['mkl']['library_dirs'].split(os.pathsep))
        # numpy's site.cfg splits libraries by comma, but numexpr historically split by os.pathsep.
        # For compatibility, we split by both.
        libs.extend(
            site['mkl']['libraries'].replace(os.pathsep, ',').split(','))
        def_macros.append(('USE_VML', None))
        print(f'FOUND MKL IMPORT')
        

def setup_package():

    parse_site_cfg()

    numexpr_extension = Extension(
        'numexpr.interpreter',
        include_dirs=inc_dirs,
        define_macros=def_macros,
        sources=sources,
        library_dirs=lib_dirs,
        libraries=libs,
        extra_compile_args=extra_cflags,
        extra_link_args=extra_link_args,
    )

    metadata = dict(
        install_requires=requirements,
        libraries=clibs,
        ext_modules=[
            numexpr_extension
        ],
    )
    setup(**metadata)


if __name__ == '__main__':
    setup_package()