paddle_cluster/setup.py at paddle · PFCCLab/paddle_cluster · 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
import subprocess
import sys

import paddle
from setuptools import find_packages
from setuptools import setup
from setuptools.command.develop import develop
from setuptools.command.install import install

from setup_ops import get_version

__version__ = get_version()


install_requires = [
    "scipy",
]

test_requires = [
    "pytest",
    "pytest-cov",
]

support_commands = ("install", "develop")
assert (
    sys.argv[1] in support_commands
), f"paddle-cluster setup.py only support {support_commands} commands."


class CustomCommand:
    """Common functionality for install and develop commands"""

    def run_setup_ops_install(self):
        try:
            subprocess.check_call([sys.executable, "setup_ops.py", "install"])
        except subprocess.CalledProcessError as e:
            print(f"Error running setup_ops.py: {e}")
            raise
        except Exception as e:
            print(f"Unexpected error: {e}")
            raise


class CustomInstallCommand(install, CustomCommand):
    def run(self):
        self.run_setup_ops_install()
        install.run(self)


class CustomDevelopCommand(develop, CustomCommand):
    def run(self):
        self.run_setup_ops_install()
        develop.run(self)


include_package_data = True
if paddle.device.cuda.device_count() > 0:
    include_package_data = False


if __name__ == "__main__":
    setup(
        name="paddle_cluster",
        version=__version__,
        description=(
            "Paddle Extension Library of Optimized Graph Cluster "
            "Algorithms. Originally from https://github.com/rusty1s/pytorch_cluster."
        ),
        author="Ruibin Cheung",
        author_email="beinggod@foxmail.com",
        keywords=[
            "paddlepaddle",
            "geometric-deep-learning",
            "graph-neural-networks",
            "cluster-algorithms",
        ],
        python_requires=">=3.8",
        install_requires=install_requires,
        extras_require={
            "test": test_requires,
        },
        cmdclass={"install": CustomInstallCommand, "develop": CustomDevelopCommand},
        packages=find_packages(),
        include_package_data=include_package_data,
    )