setup.py

# -*- coding: utf-8 -*-
# MegFlow is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2019-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

#!/usr/bin/env python
# coding=utf-8
import sys
import os
import glob
import re

from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py
from distutils.file_util import copy_file
from distutils.dir_util import copy_tree, mkpath, remove_tree
import subprocess as sp

import platform
system = platform.system().lower()

dyn_ext = 'so'
if system == 'darwin':
    dyn_ext = 'dylib'
elif system == 'windows':
    dyn_ext = 'dll'

debug = os.environ.get("DEBUG")
target_dir = os.environ.get("CARGO_TARGET_DIR")

if not debug:
    debug = False
if not target_dir:
    target_dir = "../target"


class FirstBuildExt(build_py):
    def run(self):
        self.run_command("build_ext")
        return super().run()


class CargoExtension(Extension):
    def __init__(self,
                 target,
                 src,
                 dst,
                features=[]):
        Extension.__init__(self, target, sources=[])
        self.target = target
        self.src = src
        self.dst = dst
        self.features = features

    def build(self):
        command = ['cargo', 'build', '-p', self.target]
        if len(self.features) != 0:
            command.append('--features')
            command.append(' '.join(self.features))
        if not debug:
            command.append('--release')
        sp.check_call(command)

    def install(self, prefix):
        copy_file('{}/{}'.format(prefix, self.src), 'megflow/{}'.format(self.dst))


class CopyExtension(Extension):
    def __init__(self, pattern, src, dst):
        Extension.__init__(self, '', sources=[])
        self.src = src
        self.dst = dst
        self.pattern = pattern

    def copy(self):
        mkpath('megflow/{}'.format(self.dst))
        paths = glob.glob(self.src)
        paths = [ x for x in paths if self.pattern.fullmatch(x) ]

        for path in paths:
            copy_file(path, 'megflow/{}'.format(self.dst))


class ExtBuild(build_ext):
    def run(self):
        current_dir = os.getcwd()
        repo = os.path.dirname(current_dir)
        
        prefix = target_dir
        if debug:
            prefix += '/debug'
        else:
            prefix += '/release'

        for ext in self.extensions:
            if isinstance(ext, CargoExtension):
                ext.build()
                ext.install(prefix)
            if isinstance(ext, CopyExtension):
                ext.copy()


if __name__ == '__main__':
    ext_modules=[
        CargoExtension("flow-python", f"libflow_python.{dyn_ext}", f"megflow.{dyn_ext}", features=["extension-module"]), 
        CargoExtension("flow-quickstart", "megflow_quickstart", "megflow_quickstart_inner"),
    ]

    ffmpeg_dir = os.getenv('FFMPEG_DIR')
    prebuild = os.getenv('CARGO_FEATURE_DYNAMIC')
    if prebuild is not None and ffmpeg_dir is not None:
        pattern = re.compile(f'.*?{dyn_ext}\.[0-9]*')
        ext_modules.append(CopyExtension(pattern, f"{ffmpeg_dir}/lib/*.{dyn_ext}.*", "lib/"))

    current_dir = os.getcwd()
    with open(current_dir+'/Cargo.toml') as f:
        pattern = re.compile(r'\d+\.(?:\d+\.)*\d+')
        for line in f:
            if line.startswith('version'):
                version = re.search(pattern, line).group()
                break

    setup(
        options={
            'bdist_wheel': {
                'py_limited_api': "cp36",
            }
        },
        name="megflow",
        version=version,
        packages=["megflow"],
        author="Megvii IPU-SDK Team",
        author_email="megengine@megvii.com",
        url="https://github.com/MegEngine/MegFlow",
        include_package_data=True,
        classifiers=[
            'Development Status :: 3 - Alpha',
            'Intended Audience :: Developers',
            'License :: OSI Approved :: Apache Software License',
            'Natural Language :: English',
            'Operating System :: POSIX :: Linux',
            'Programming Language :: Rust',
            'Programming Language :: Python :: 3',
            'Topic :: Software Development :: Libraries :: Application Frameworks',
            'Topic :: Scientific/Engineering',
            'Topic :: Scientific/Engineering :: Mathematics',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
            'Topic :: Software Development',
            'Topic :: Software Development :: Libraries',
            'Topic :: Software Development :: Libraries :: Python Modules',
        ],
        ext_modules=ext_modules,
        package_data={"": [f'megflow.{dyn_ext}', 'lib/*', 'megflow_quickstart_inner']},
        entry_points={
            'console_scripts':['megflow_run=megflow.command_line:megflow_run', 'run_with_plugins=megflow.command_line:run_with_plugins', 'megflow_quickstart=megflow.command_line:megflow_quickstart'],
        },
        cmdclass={
            'build_ext': ExtBuild,
            'build_py': FirstBuildExt
        },
        zip_safe=False,
    )