py2so.py 6.15 KB
import os
import sys
import time
import shutil
import argparse
from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
from Cython.Distutils import build_ext


if sys.version_info[0] < 3:
    raise Exception("This program requires at least python3.6")


parser = argparse.ArgumentParser(description='compile the .py to .so(Linux/Mac) or .pdy(Win)')
parser.add_argument('-i', '--input', help='the directory(file) path of your project', required=True)
parser.add_argument('-o', '--output', help='the directory path of compiled file')
parser.add_argument('-e', '--exclude', help='exclude file. eg: ignore/ignore.py,main.py,ignore')
args = parser.parse_args()


def get_exclude_paths(input_path):
    exclude_paths = []
    if args.exclude is not None and os.path.isdir(input_path):
        paths = args.exclude.split(',')
        for path in paths:
            abs_path = path if os.path.isabs(path) else os.path.join(input_path, path)
            if not os.path.exists(abs_path):
                print('Warning: exclude path not exists: {0}'.format(abs_path))
                continue
            exclude_paths.append(abs_path)
    return set(exclude_paths)


def get_py_files(dir_path, exclude_paths):
    for parent, _, filenames in os.walk(dir_path):
        if parent in exclude_paths:
            print('exclude path: {0}'.format(parent))
            continue
        for filename in filenames:
            if not filename.endswith('.py'):
                continue
            file_path = os.path.join(parent, filename)
            if file_path in exclude_paths:
                print('exclude path: {0}'.format(file_path))
                continue
            yield file_path


def is_about_py(filename):
    if filename.endswith('.py') or filename.endswith('.pyc') or filename.endswith('.pyo') or filename.endswith('.pyd'):
        return True
    return False


def copy_other_files(input_path, exclude_paths, output_path):
    if not os.path.isdir(input_path):
        return
    for parent, _, filenames in os.walk(input_path):
        if parent in exclude_paths:
            relpath = os.path.relpath(parent, input_path)
            target_dir = os.path.join(output_path, relpath)
            shutil.copytree(parent, target_dir)
            print('copy {0} to {1}'.format(parent, target_dir))
            continue
        for filename in filenames:
            file_path = os.path.join(parent, filename)
            if file_path in exclude_paths or not is_about_py(filename):
                relpath = os.path.relpath(file_path, input_path)
                target_path = os.path.join(output_path, relpath)
                os.makedirs(os.path.dirname(target_path), exist_ok=True)
                shutil.copy(file_path, target_path)
                print('copy {0} to {1}'.format(file_path, target_path))


def del_c(dir_path, timestamp):
    if not os.path.isdir(dir_path):
        c_file_path = '{0}.c'.format(os.path.splitext(dir_path)[0])
        os.remove(c_file_path)
        print('remove c file path: {0}'.format(c_file_path))
        return
    for parent, _, filenames in os.walk(dir_path):
        for filename in filenames:
            file_path = os.path.join(parent, filename)
            if filename.endswith('.c') and os.stat(file_path).st_mtime > timestamp:
                os.remove(file_path)
                print('remove c file path: {0}'.format(file_path))


def get_extensions(file_paths, input_path):
    print('---------- python file summary ----------')
    extensions = []
    for file_path in file_paths:
        if os.path.isfile(input_path):
            relpath = os.path.basename(file_path)
        else:
            relpath = os.path.relpath(file_path, start=input_path)
        name = os.path.splitext(relpath)[0].replace(os.path.sep, '.')
        sources = [file_path]
        print('python file path: {0} module name: {1}'.format(file_path, name))
        extension = Extension(name, sources)
        extension.cython_c_in_temp = True
        extensions.append(extension)
    return extensions


def main():
    input_path = os.path.realpath(args.input)
    if not os.path.exists(input_path):
        print('input path not exists: {0}'.format(input_path))
        return

    timestamp = int(time.time())
    build_dir_name = 'build_{0}'.format(timestamp)
    tmp_dir_name = 'tmp_{0}'.format(timestamp)
    if args.output is None:
        output_path = os.path.join(os.path.dirname(input_path), build_dir_name)
    else:
        output_path = os.path.realpath(args.output)
    tmp_path = os.path.join(output_path, tmp_dir_name)
    os.makedirs(tmp_path, exist_ok=True)

    exclude_paths = get_exclude_paths(input_path)

    if os.path.isdir(input_path):
        file_paths = get_py_files(input_path, exclude_paths)
    else:
        file_paths = [input_path]

    print('---------- directory summary ----------')
    print('input path: {0}'.format(input_path))
    print('output path: {0}'.format(output_path))
    print('output tmp path: {0}'.format(tmp_path))
    print('exclude paths: {0}'.format(exclude_paths))
    extensions = get_extensions(file_paths, input_path)

    print('---------- compile start ----------')
    success = False
    try:
        compiler_directives = {"language_level": '3'}
        ext_modules = cythonize(module_list=extensions,
                                compiler_directives=compiler_directives)

        setup(cmdclass={'build_ext': build_ext},
              ext_modules=ext_modules, script_args=["build_ext", "-b", output_path, "-t", tmp_path])
    except Extension as e:
        print('---------- compile failed ----------')
        print('compile error: {0}'.format(e))
    else:
        print('---------- compile completed ----------')
        print('output path: {0}'.format(output_path))
        success = True
    finally:
        print('---------- cleaning ----------')
        if os.path.exists(tmp_path):
            shutil.rmtree(tmp_path)
            print('remove output tmp path: {0}'.format(tmp_path))
        del_c(input_path, timestamp)

    if success:
        print('---------- copy !.py file to {0} ----------'.format(build_dir_name))
        copy_other_files(input_path, exclude_paths, output_path)


if __name__ == '__main__':
    main()