https://github.com/google-deepmind/tree/setup.py
"""Setup for pip package."""
import os, platform, sys, sysconfig, shutil, subprocess, setuptools
from setuptools.command import build_ext
here = os.path.dirname(os.path.abspath(file))
def _get_tree_version():
"""Parse the version string from tree/__init__.py."""
with open(os.path.join(here, 'tree', '__init__.py')) as f:
try:
version_line = next(line for line in f if line.startswith('__version__'))
except StopIteration:
raise ValueError('__version__ not defined in tree/__init__.py')
else:
ns = {}
exec(version_line, ns) # pylint: disable=exec-used
return ns['__version__']
def _parse_requirements(path):
with open(os.path.join(here, path)) as f:
return [
line.rstrip() for line in f
if not (line.isspace() or line.startswith('#'))
]
class CMakeExtension(setuptools.Extension):
"""An extension with no sources.
We do not want distutils to handle any of the compilation (instead we rely
on CMake), so we always pass an empty list to the constructor.
"""
def __init__(self, name, source_dir=''):
super().__init__(name, sources=[])
self.source_dir = os.path.abspath(source_dir)
class BuildCMakeExtension(build_ext.build_ext):
"""Our custom build_ext command.
Uses CMake to build extensions instead of a bare compiler (e.g. gcc, clang).
"""
def run(self):
self._check_build_environment()
for ext in self.extensions:
self.build_extension(ext)
def _check_build_environment(self):
"""Check for required build tools: CMake, C++ compiler, and python dev."""
try:
subprocess.check_call(['cmake', '--version'])
except OSError as e:
ext_names = ', '.join(e.name for e in self.extensions)
raise RuntimeError(
f'CMake must be installed to build the following extensions: {ext_names}'
) from e
print('Found CMake')
def build_extension(self, ext):
extension_dir = os.path.abspath(
os.path.dirname(self.get_ext_fullpath(ext.name)))
build_cfg = 'Debug' if self.debug else 'Release'
cmake_args = [
f'-DPython3_ROOT_DIR={sys.prefix}',
f'-DPython3_EXECUTABLE={sys.executable}',
f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extension_dir}',
f'-DCMAKE_BUILD_TYPE={build_cfg}'
]
if platform.system() != 'Windows':
cmake_args.extend([
f'-DPython3_LIBRARY={sysconfig.get_paths()["stdlib"]}',
f'-DPython3_INCLUDE_DIR={sysconfig.get_paths()["include"]}',
])
if platform.system() == 'Darwin' and os.environ.get('ARCHFLAGS'):
osx_archs = []
if '-arch x86_64' in os.environ['ARCHFLAGS']:
osx_archs.append('x86_64')
if '-arch arm64' in os.environ['ARCHFLAGS']:
osx_archs.append('arm64')
cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={";".join(osx_archs)}')
os.makedirs(self.build_temp, exist_ok=True)
subprocess.check_call(
['cmake', ext.source_dir] + cmake_args, cwd=self.build_temp)
subprocess.check_call(
['cmake', '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
cwd=self.build_temp)
# Force output to <extension_dir>/. Amends CMake multigenerator output paths
# on Windows and avoids Debug/ and Release/ subdirs, which is CMake default.
tree_dir = os.path.join(extension_dir, 'tree') # pylint:disable=unreachable
for cfg in ('Release', 'Debug'):
cfg_dir = os.path.join(extension_dir, cfg)
if os.path.isdir(cfg_dir):
for f in os.listdir(cfg_dir):
shutil.move(os.path.join(cfg_dir, f), tree_dir)
setuptools.setup(
name='dm-tree',
version=_get_tree_version(),
url='https://github.com/deepmind/tree',
description='Tree is a library for working with nested data structures.',
author='DeepMind',
author_email='tree-copybara@google.com',
long_description=open(os.path.join(here, 'README.md')).read(),
long_description_content_type='text/markdown',
# Contained modules and scripts.
packages=setuptools.find_packages(),
tests_require=_parse_requirements('requirements-test.txt'),
test_suite='tree',
cmdclass=dict(build_ext=BuildCMakeExtension),
ext_modules=[CMakeExtension('_tree', source_dir='tree')],
zip_safe=False,
# PyPI package information.
classifiers=[
'Development Status :: 4 - Beta',
'Intended Audience :: Developers',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: Apache Software License',
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'Programming Language :: Python :: 3.9',
'Programming Language :: Python :: 3.10',
'Programming Language :: Python :: 3.11',
'Topic :: Scientific/Engineering :: Mathematics',
'Topic :: Software Development :: Libraries',
],
license='Apache 2.0',
keywords='tree nest flatten',
)