1# Support code for building a C extension with zstd files
2#
3# Copyright (c) 2016-present, Gregory Szorc
4#               2017-present, Thomas Waldmann (mods to make it more generic)
5# All rights reserved.
6#
7# This software may be modified and distributed under the terms
8# of the BSD license. See the LICENSE file for details.
9
10import os
11
12# zstd files, structure as seen in zstd project repository:
13
14zstd_sources = [
15    'lib/common/debug.c',
16    'lib/common/entropy_common.c',
17    'lib/common/error_private.c',
18    'lib/common/fse_decompress.c',
19    'lib/common/pool.c',
20    'lib/common/threading.c',
21    'lib/common/xxhash.c',
22    'lib/common/zstd_common.c',
23    'lib/compress/fse_compress.c',
24    'lib/compress/hist.c',
25    'lib/compress/huf_compress.c',
26    'lib/compress/zstd_compress.c',
27    'lib/compress/zstd_compress_literals.c',
28    'lib/compress/zstd_compress_sequences.c',
29    'lib/compress/zstd_compress_superblock.c',
30    'lib/compress/zstd_double_fast.c',
31    'lib/compress/zstd_fast.c',
32    'lib/compress/zstd_lazy.c',
33    'lib/compress/zstd_ldm.c',
34    'lib/compress/zstd_opt.c',
35    'lib/compress/zstdmt_compress.c',
36    'lib/decompress/huf_decompress.c',
37    'lib/decompress/zstd_ddict.c',
38    'lib/decompress/zstd_decompress.c',
39    'lib/decompress/zstd_decompress_block.c',
40    'lib/dictBuilder/cover.c',
41    'lib/dictBuilder/divsufsort.c',
42    'lib/dictBuilder/fastcover.c',
43    'lib/dictBuilder/zdict.c',
44]
45
46zstd_sources_legacy = [
47    'lib/deprecated/zbuff_common.c',
48    'lib/deprecated/zbuff_compress.c',
49    'lib/deprecated/zbuff_decompress.c',
50    'lib/legacy/zstd_v01.c',
51    'lib/legacy/zstd_v02.c',
52    'lib/legacy/zstd_v03.c',
53    'lib/legacy/zstd_v04.c',
54    'lib/legacy/zstd_v05.c',
55    'lib/legacy/zstd_v06.c',
56    'lib/legacy/zstd_v07.c',
57]
58
59zstd_includes = [
60    'lib',
61    'lib/common',
62    'lib/compress',
63    'lib/decompress',
64    'lib/dictBuilder',
65]
66
67zstd_includes_legacy = [
68    'lib/deprecated',
69    'lib/legacy',
70]
71
72
73def zstd_system_prefix(prefixes):
74    for prefix in prefixes:
75        filename = os.path.join(prefix, 'include', 'zstd.h')
76        if os.path.exists(filename):
77            with open(filename, 'rb') as fd:
78                if b'ZSTD_getFrameContentSize' in fd.read():  # checks for zstd >= 1.3.0
79                    return prefix
80
81
82def zstd_ext_kwargs(bundled_path, system_prefix=None, system=False, multithreaded=False, legacy=False, **kwargs):
83    """amend kwargs with zstd suff for a distutils.extension.Extension initialization.
84
85    bundled_path: relative (to this file) path to the bundled library source code files
86    system_prefix: where the system-installed library can be found
87    system: True: use the system-installed shared library, False: use the bundled library code
88    multithreaded: True: define ZSTD_MULTITHREAD
89    legacy: include legacy API support
90    kwargs: distutils.extension.Extension kwargs that should be amended
91    returns: amended kwargs
92    """
93    def multi_join(paths, *path_segments):
94        """apply os.path.join on a list of paths"""
95        return [os.path.join(*(path_segments + (path, ))) for path in paths]
96
97    use_system = system and system_prefix is not None
98
99    sources = kwargs.get('sources', [])
100    if not use_system:
101        sources += multi_join(zstd_sources, bundled_path)
102        if legacy:
103            sources += multi_join(zstd_sources_legacy, bundled_path)
104
105    include_dirs = kwargs.get('include_dirs', [])
106    if use_system:
107        include_dirs += multi_join(['include'], system_prefix)
108    else:
109        include_dirs += multi_join(zstd_includes, bundled_path)
110        if legacy:
111            include_dirs += multi_join(zstd_includes_legacy, bundled_path)
112
113    library_dirs = kwargs.get('library_dirs', [])
114    if use_system:
115        library_dirs += multi_join(['lib'], system_prefix)
116
117    libraries = kwargs.get('libraries', [])
118    if use_system:
119        libraries += ['zstd', ]
120
121    extra_compile_args = kwargs.get('extra_compile_args', [])
122    if multithreaded:
123        extra_compile_args += ['-DZSTD_MULTITHREAD', ]
124    if not use_system:
125        extra_compile_args += ['-DZSTDLIB_VISIBILITY=', '-DZDICTLIB_VISIBILITY=', '-DZSTDERRORLIB_VISIBILITY=', ]
126                               # '-fvisibility=hidden' does not work, doesn't find PyInit_compress then
127        if legacy:
128            extra_compile_args += ['-DZSTD_LEGACY_SUPPORT=1', ]
129
130    ret = dict(**kwargs)
131    ret.update(dict(sources=sources, extra_compile_args=extra_compile_args,
132                    include_dirs=include_dirs, library_dirs=library_dirs, libraries=libraries))
133    return ret
134