backports.zstd/setup.py at master · Rogdham/backports.zstd · 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
import platform
import sys
import sysconfig
from distutils.unixccompiler import UnixCCompiler
from pathlib import Path

from setuptools import Extension, setup

if not ((3, 9) <= sys.version_info < (3, 14)):
    raise RuntimeError(f"Unsupported Python version: {sys.version}")


# create a LICENSE_zstd.txt file
# wheels distributions needs to ship the license of the zstd library
ROOT_PATH = Path(__file__).parent.absolute()
with (ROOT_PATH / "LICENSE_zstd.txt").open("w") as f:
    f.write(
        "Depending on how it is built, this package may distribute the zstd library,\n"
        "partially or in its integrality, in source or binary form.\n\n"
        "Its license is reproduced below.\n\n"
        "---\n\n"
    )
    f.write((ROOT_PATH / "src" / "c" / "zstd" / "LICENSE").read_text())


UnixCCompiler.src_extensions.append(".S")


_PLATFORM_IS_WIN = sysconfig.get_platform().startswith("win")
_USE_CFFI = platform.python_implementation() == "PyPy"
try:
    sys.argv.remove("--system-zstd")
except ValueError:
    _SYSTEM_ZSTD = False
else:
    _SYSTEM_ZSTD = True


def locate_sources(*sub_paths):
    extensions = "cC" if _PLATFORM_IS_WIN else "cCsS"
    yield from map(str, Path(*sub_paths).rglob(f"*.[{extensions}]"))


def build_extension():
    kwargs = dict(
        sources=[],
        include_dirs=[],
        libraries=[],
        extra_compile_args=[],
        extra_link_args=[],
        define_macros=[
            ("ZSTD_MULTITHREAD", None),  # enable multithreading support
        ],
    )

    if _PLATFORM_IS_WIN:
        kwargs["extra_compile_args"] += ["/Ob3", "/GF", "/Gy"]
    else:
        kwargs["extra_compile_args"] += ["-g0", "-flto"]
        kwargs["extra_link_args"] += ["-g0", "-flto"]

    if _SYSTEM_ZSTD:
        kwargs["libraries"].append("zstd")
    else:
        kwargs["sources"] += [
            *locate_sources("src", "c", "zstd", "lib", "common"),
            *locate_sources("src", "c", "zstd", "lib", "compress"),
            *locate_sources("src", "c", "zstd", "lib", "decompress"),
            *locate_sources("src", "c", "zstd", "lib", "dictBuilder"),
        ]
        kwargs["include_dirs"] += [
            "src/c/zstd/lib",
            "src/c/zstd/lib/common",
            "src/c/zstd/lib/dictBuilder",
        ]

    if _USE_CFFI:
        import cffi

        ffibuilder = cffi.FFI()
        ffibuilder.cdef((ROOT_PATH / "src" / "c" / "cffi" / "cdef.h").read_text())
        ffibuilder.set_source(
            source=(ROOT_PATH / "src" / "c" / "cffi" / "source.c").read_text(),
            module_name="backports.zstd._zstd_cffi",
            **kwargs,
        )
        return ffibuilder.distutils_extension()

    else:
        kwargs["sources"] += [
            *locate_sources("src", "c", "compression_zstd"),
            *locate_sources("src", "c", "compat"),
        ]
        kwargs["include_dirs"] += [
            "src/c/compat",
            "src/c/compression_zstd",
            "src/c/compression_zstd/clinic",
            "src/c/pythoncapi-compat",
        ]
        return Extension(
            name="backports.zstd._zstd",
            **kwargs,
        )


setup(ext_modules=[build_extension()])