setup.py

"""
Setup script for KortexDL Python bindings

This script uses CMake to build the C++ extension module.
Requires: Intel oneAPI (MKL), pybind11, CMake, and a C++ compiler (icpx preferred)

Usage:
    # Development install (editable)
    pip install -e .
    
    # Build wheel
    pip wheel .
"""

import os
import sys
import subprocess
import shutil
import glob
from pathlib import Path

from setuptools import setup, Extension, find_packages
from setuptools.command.build_ext import build_ext


class CMakeExtension(Extension):
    """Custom extension class for CMake-based builds"""
    def __init__(self, name, sourcedir=""):
        super().__init__(name, sources=[])
        self.sourcedir = os.path.abspath(sourcedir)


class CMakeBuild(build_ext):
    """Custom build_ext command that runs CMake"""
    
    def run(self):
        # Check for CMake
        try:
            subprocess.check_output(["cmake", "--version"])
        except OSError:
            raise RuntimeError("CMake must be installed to build this extension")
        
        for ext in self.extensions:
            self.build_extension(ext)
    
    def build_extension(self, ext):
        # Get the extension output path
        ext_fullpath = Path(self.get_ext_fullpath(ext.name)).resolve()
        extdir = ext_fullpath.parent
        
        # Ensure output directory exists
        extdir.mkdir(parents=True, exist_ok=True)
        
        # CMake configuration
        cmake_args = [
            f"-DPYTHON_EXECUTABLE={sys.executable}",
            "-DCMAKE_BUILD_TYPE=Release",
        ]
        
        # Try to use Intel compiler if available
        icpx = shutil.which("icpx")
        icx = shutil.which("icx")
        if icpx and icx:
            cmake_args.extend([
                f"-DCMAKE_CXX_COMPILER={icpx}",
                f"-DCMAKE_C_COMPILER={icx}",
            ])
            print(f"Using Intel compiler: {icpx}")
        else:
            print("Warning: Intel compiler not found, using default compiler")
        
        # Build arguments
        build_args = ["--config", "Release", "--target", "_kortexdl_core", "-j"]
        
        # Create build directory in source tree (for finding kortexdl-cpp)
        build_temp = Path(ext.sourcedir) / "build"
        build_temp.mkdir(parents=True, exist_ok=True)
        
        # Run CMake configure
        print(f"CMake configure in {build_temp}")
        print(f"  Source: {ext.sourcedir}")
        print(f"  Output: {extdir}")
        
        try:
            subprocess.check_call(
                ["cmake", ext.sourcedir] + cmake_args,
                cwd=build_temp,
            )
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"CMake configuration failed: {e}")
        
        # Run CMake build
        print("CMake build...")
        try:
            subprocess.check_call(
                ["cmake", "--build", "."] + build_args,
                cwd=build_temp,
            )
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"CMake build failed: {e}")
        
        # Find and copy the built .so file
        # CMakeLists.txt outputs to source directory
        source_so_patterns = [
            Path(ext.sourcedir) / f"_kortexdl_core*.so",
            Path(ext.sourcedir) / f"_kortexdl_core*.pyd",
        ]
        
        built_files = []
        for pattern in source_so_patterns:
            built_files.extend(glob.glob(str(pattern)))
        
        if not built_files:
            # Check build directory too
            build_patterns = [
                build_temp / f"_kortexdl_core*.so",
                build_temp / f"_kortexdl_core*.pyd",
            ]
            for pattern in build_patterns:
                built_files.extend(glob.glob(str(pattern)))
        
        if built_files:
            src_file = Path(built_files[0])
            dst_file = ext_fullpath
            # Skip if same file (editable install)
            if src_file.resolve() != dst_file.resolve():
                print(f"Copying {src_file} -> {dst_file}")
                shutil.copy2(src_file, dst_file)
            print(f"✅ Extension installed successfully!")
        else:
            # If already built in correct location, that's fine
            if ext_fullpath.exists():
                print(f"✅ Extension already at {ext_fullpath}")
            else:
                raise RuntimeError(f"Built extension not found! Searched: {source_so_patterns}")


# Read README
readme_file = Path(__file__).parent / "README.md"
long_description = readme_file.read_text() if readme_file.exists() else ""

# Check for pre-built binary
prebuilt_so = list(Path(__file__).parent.glob("_kortexdl_core*.so"))

setup(
    name="kortexdl",
    version="2.0.0",
    author="Mostafa Saad",
    author_email="",
    description="High-performance deep learning framework with Intel oneAPI MKL",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/Mostafasaad1/kortexdl-python",
    
    # Python module
    py_modules=["kortexdl"],
    
    # C++ extension
    ext_modules=[CMakeExtension("_kortexdl_core")],
    cmdclass={"build_ext": CMakeBuild},
    
    # Requirements
    python_requires=">=3.8",
    install_requires=[
        "numpy>=1.20.0",
    ],
    extras_require={
        "dev": [
            "pytest>=6.0",
            "pytest-cov>=2.0",
            "pybind11>=2.10",
            "black>=22.0",
            "mypy>=0.950",
            "jupyter",
            "matplotlib",
        ],
        "test": [
            "pytest>=6.0",
            "pytest-cov>=2.0",
        ],
        "notebooks": [
            "jupyter",
            "matplotlib",
        ],
    },
    
    # Include pre-built .so if available
    package_data={"": ["*.so", "*.pyd", "*.dll"]},
    include_package_data=True,
    
    # Metadata
    classifiers=[
        "Development Status :: 4 - Beta",
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "License :: OSI Approved :: MIT License",
        "Programming Language :: C++",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: Python :: 3.12",
    ],
    
    zip_safe=False,
)