import os
import re
import pathlib
import platform
from distutils.file_util import copy_file
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext as _build_ext
class PrecompiledExtesion(Extension):
def __init__(self, name):
super().__init__(name, sources=[])
class build_ext(_build_ext):
def build_extension(self, ext):
if not isinstance(ext, PrecompiledExtesion):
return super().build_extension(ext)
if not self.inplace:
fullpath = self.get_ext_fullpath(ext.name)
extdir = pathlib.Path(fullpath)
extdir.parent.mkdir(parents=True, exist_ok=True)
modpath = self.get_ext_fullname(ext.name).split('.')
if platform.system() == 'Windows':
modpath[-1] += '.pyd'
else:
modpath[-1] += '.so'
modpath = str(pathlib.Path(*modpath).resolve())
copy_file(modpath, fullpath, verbose=self.verbose, dry_run=self.dry_run)
package_name = 'MegEngine'
v = {}
with open("megengine/version.py") as fp:
exec(fp.read(), v)
__version__ = v['__version__']
email = 'megengine@megvii.com'
public_version_postfix = os.environ.get('PUBLIC_VERSION_POSTFIX')
if public_version_postfix:
__version__ = '{}{}'.format(__version__, public_version_postfix)
local_version = []
strip_sdk_info = os.environ.get('STRIP_SDK_INFO', 'False').lower()
sdk_name = os.environ.get('SDK_NAME', 'cpu')
if 'true' == strip_sdk_info:
print('wheel version strip sdk info')
else:
local_version.append(sdk_name)
local_postfix = os.environ.get('LOCAL_VERSION')
if local_postfix:
local_version.append(local_postfix)
if len(local_version):
__version__ = '{}+{}'.format(__version__, '.'.join(local_version))
packages = find_packages(exclude=['test'])
megengine_data = [
str(f.relative_to('megengine'))
for f in pathlib.Path('megengine', 'core', 'include').glob('**/*')
]
megengine_data += [
str(f.relative_to('megengine'))
for f in pathlib.Path('megengine', 'core', 'lib').glob('**/*')
]
megenginelite_data = [
str(f.relative_to('megenginelite'))
for f in pathlib.Path('megenginelite').glob('**/*')
]
if platform.system() == 'Windows':
megenginelite_data.remove('libs\\liblite_shared_whl.pyd')
else:
megenginelite_data.remove('libs/liblite_shared_whl.so')
with open('requires.txt') as f:
requires = f.read().splitlines()
with open('requires-style.txt') as f:
requires_style = f.read().splitlines()
with open('requires-test.txt') as f:
requires_test = f.read().splitlines()
prebuild_modules=[PrecompiledExtesion('megengine.core._imperative_rt')]
prebuild_modules.append(PrecompiledExtesion('megenginelite.libs.liblite_shared_whl'))
setup_kwargs = dict(
name=package_name,
version=__version__,
description='Framework for numerical evaluation with '
'auto-differentiation',
author='Megvii Engine Team',
author_email=email,
packages=packages,
package_data={
'megengine': megengine_data,
'megenginelite': megenginelite_data,
},
ext_modules=prebuild_modules,
install_requires=requires,
extras_require={
'dev': requires_style + requires_test,
'ci': requires_test,
},
cmdclass={'build_ext': build_ext},
scripts = ['./megengine/tools/mge'],
)
setup_kwargs.update(dict(
classifiers=[
'Development Status :: 3 - Alpha',
'Intended Audience :: Developers',
'Intended Audience :: Education',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: Apache Software License',
'Programming Language :: C++',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.5',
'Programming Language :: Python :: 3.6',
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'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',
],
license='Apache 2.0',
keywords='megengine deep learning',
data_files = [("megengine", [
"../LICENSE",
"../ACKNOWLEDGMENTS",
])]
))
setup(**setup_kwargs)