TensorComprehensions/setup.py at master · prigoyal/TensorComprehensions · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from setuptools import setup, Extension, distutils, Command, find_packages
import setuptools.command.install
import setuptools.command.develop
import setuptools.command.build_py
import distutils.command.build
import subprocess, os

TOP_DIR = os.path.realpath(os.path.dirname(__file__))

print('TOP_DIR: ', TOP_DIR)

###############################################################################
# Version and build number
###############################################################################

try:
    git_version = subprocess.check_output(
        ['git', 'rev-parse', 'HEAD'], cwd=TOP_DIR
    ).decode('ascii').strip()
except (OSError, subprocess.CalledProcessError):
    git_version = None

tc_version = '0.1.0'
tc_build_number = 1
# versioning should comply with
# https://www.python.org/dev/peps/pep-0440/#public-version-identifiers
# when conda packaging, we get these values from environment variables
if os.getenv('TC_BUILD_VERSION') and os.getenv('TC_BUILD_NUMBER'):
    assert os.getenv('TC_BUILD_NUMBER') is not None, "Please specify valid build number"
    tc_build_number = int(os.getenv('TC_BUILD_NUMBER'))
    tc_version = str(os.getenv('TC_BUILD_VERSION'))
    if tc_build_number > 1:
        tc_version += '.post' + str(tc_build_number)

print('git_version: {} tc_version: {} tc_build_number: {}'.format(
    git_version, tc_version, tc_build_number
))

################################################################################
# Custom override commands
################################################################################

class build_py(setuptools.command.build_py.build_py):
    def run(self):
        self.create_version_file()
        setuptools.command.build_py.build_py.run(self)

    @staticmethod
    def create_version_file():
        global tc_version, tc_build_number, git_version
        print('BUILD git_version: {} tc_version: {} tc_build_number: {}'.format(
            git_version, tc_version, tc_build_number
        ))
        with open(
            os.path.join(TOP_DIR, 'tensor_comprehensions', 'version.py'), 'w'
        ) as fopen:
            fopen.write("__version__ = '{}'\n".format(str(tc_version)))
            fopen.write("build_number = {}\n".format(int(tc_build_number)))
            fopen.write("git_version = '{}'\n".format(str(git_version)))
        print('Version file written.')


class install(setuptools.command.install.install):
    def run(self):
        setuptools.command.install.install.run(self)


class develop(setuptools.command.develop.develop):
    def run(self):
        build_py.create_version_file()
        setuptools.command.develop.develop.run(self)

################################################################################
# Extensions
################################################################################

# we have build extensions in the cmake command so no extensions left to build
print('All extension module were build with cmake')
ext_modules = []

################################################################################
# Command line options
################################################################################

cmdclass = {
    'develop': develop,
    'build_py': build_py,
    'install': install,
}

###############################################################################
# Pure python packages
###############################################################################
# don't include the tests in the conda package, we run these tests when building
# the package
packages = find_packages(exclude=('test_python', 'test_python.*'))

###############################################################################
# Main
###############################################################################


setup(
    name="tensor_comprehensions",
    version=tc_version,
    ext_modules=ext_modules,
    cmdclass=cmdclass,
    packages=packages,
    package_data={'tensor_comprehensions': [
        '*.so',
    ]},
    install_requires=['pyyaml', 'numpy'],
    author='prigoyal',
    author_email='prigoyal@fb.com',
    url='https://github.com/nicolasvasilache/c2isl/',
    license="Apache 2.0",
    description=("Framework-Agnostic Abstractions for High-Performance Machine Learning"),
)