|
| 1 | +""" |
| 2 | +Hatchling custom build hook for CFFI extension modules. |
| 3 | +
|
| 4 | +This builds the NVX (Native Vector Extensions) for WebSocket frame masking |
| 5 | +and UTF-8 validation using CFFI. |
| 6 | +
|
| 7 | +See: https://hatch.pypa.io/latest/plugins/build-hook/custom/ |
| 8 | +""" |
| 9 | + |
| 10 | +import os |
| 11 | +import sys |
| 12 | +import sysconfig |
| 13 | +import importlib.util |
| 14 | +from pathlib import Path |
| 15 | + |
| 16 | +from hatchling.builders.hooks.plugin.interface import BuildHookInterface |
| 17 | + |
| 18 | + |
| 19 | +class CFfiBuildHook(BuildHookInterface): |
| 20 | + """Build hook for compiling CFFI extension modules.""" |
| 21 | + |
| 22 | + PLUGIN_NAME = "cffi" |
| 23 | + |
| 24 | + def initialize(self, version, build_data): |
| 25 | + """ |
| 26 | + Called before each build. |
| 27 | +
|
| 28 | + For wheel builds, compile the CFFI modules. |
| 29 | + For sdist builds, just ensure source files are included. |
| 30 | + """ |
| 31 | + if self.target_name != "wheel": |
| 32 | + # Only compile for wheel builds, sdist just includes source |
| 33 | + return |
| 34 | + |
| 35 | + # Check if NVX build is disabled |
| 36 | + if os.environ.get("AUTOBAHN_USE_NVX", "1") in ("0", "false"): |
| 37 | + print("AUTOBAHN_USE_NVX is disabled, skipping CFFI build") |
| 38 | + return |
| 39 | + |
| 40 | + # Build CFFI modules |
| 41 | + built_extensions = self._build_cffi_modules(build_data) |
| 42 | + |
| 43 | + # If we built extensions, mark this as a platform-specific wheel |
| 44 | + if built_extensions: |
| 45 | + # Setting infer_tag tells hatchling to use platform-specific wheel tag |
| 46 | + build_data["infer_tag"] = True |
| 47 | + # Mark as having binary extensions |
| 48 | + build_data["pure_python"] = False |
| 49 | + |
| 50 | + def _get_ext_suffix(self): |
| 51 | + """Get the extension suffix for the current Python interpreter. |
| 52 | +
|
| 53 | + E.g., '.cpython-311-x86_64-linux-gnu.so' or '.pypy311-pp73-x86_64-linux-gnu.so' |
| 54 | + """ |
| 55 | + return sysconfig.get_config_var("EXT_SUFFIX") or ".so" |
| 56 | + |
| 57 | + def _build_cffi_modules(self, build_data): |
| 58 | + """Compile the CFFI extension modules using direct file execution. |
| 59 | +
|
| 60 | + Returns True if any extensions were successfully built. |
| 61 | + """ |
| 62 | + src_path = Path(self.root) / "src" |
| 63 | + nvx_dir = src_path / "autobahn" / "nvx" |
| 64 | + built_any = False |
| 65 | + |
| 66 | + # Get the extension suffix for current Python to filter artifacts |
| 67 | + ext_suffix = self._get_ext_suffix() |
| 68 | + print(f"Building for Python with extension suffix: {ext_suffix}") |
| 69 | + |
| 70 | + # CFFI module files to build |
| 71 | + cffi_modules = [ |
| 72 | + ("_utf8validator.py", "ffi"), |
| 73 | + ("_xormasker.py", "ffi"), |
| 74 | + ] |
| 75 | + |
| 76 | + for module_file, ffi_name in cffi_modules: |
| 77 | + module_path = nvx_dir / module_file |
| 78 | + print(f"Building CFFI module: {module_path}") |
| 79 | + |
| 80 | + try: |
| 81 | + # Load the module directly from file (like CFFI's setuptools integration) |
| 82 | + # This avoids triggering package-level imports |
| 83 | + spec = importlib.util.spec_from_file_location( |
| 84 | + f"_cffi_build_{module_file}", |
| 85 | + module_path |
| 86 | + ) |
| 87 | + module = importlib.util.module_from_spec(spec) |
| 88 | + |
| 89 | + # We need to set up sys.path so the module can find _compile_args.py |
| 90 | + old_path = sys.path.copy() |
| 91 | + sys.path.insert(0, str(nvx_dir)) |
| 92 | + sys.path.insert(0, str(src_path)) |
| 93 | + |
| 94 | + try: |
| 95 | + spec.loader.exec_module(module) |
| 96 | + ffi = getattr(module, ffi_name) |
| 97 | + |
| 98 | + # Compile the CFFI module |
| 99 | + # The compiled .so/.pyd goes to the current directory by default |
| 100 | + # We want it in the nvx_dir |
| 101 | + old_cwd = os.getcwd() |
| 102 | + os.chdir(nvx_dir) |
| 103 | + try: |
| 104 | + ffi.compile(verbose=True) |
| 105 | + finally: |
| 106 | + os.chdir(old_cwd) |
| 107 | + |
| 108 | + finally: |
| 109 | + sys.path = old_path |
| 110 | + |
| 111 | + # Find the compiled artifact matching CURRENT Python and add to build_data |
| 112 | + # Only include .so files that match the current interpreter's extension suffix |
| 113 | + for artifact in nvx_dir.glob("_nvx_*" + ext_suffix): |
| 114 | + src_file = str(artifact) |
| 115 | + dest_path = f"autobahn/nvx/{artifact.name}" |
| 116 | + build_data["force_include"][src_file] = dest_path |
| 117 | + print(f" -> Added artifact: {artifact.name} -> {dest_path}") |
| 118 | + built_any = True |
| 119 | + |
| 120 | + # Handle Windows .pyd files |
| 121 | + if ext_suffix.endswith(".pyd"): |
| 122 | + for artifact in nvx_dir.glob("_nvx_*" + ext_suffix): |
| 123 | + src_file = str(artifact) |
| 124 | + dest_path = f"autobahn/nvx/{artifact.name}" |
| 125 | + build_data["force_include"][src_file] = dest_path |
| 126 | + print(f" -> Added artifact: {artifact.name} -> {dest_path}") |
| 127 | + built_any = True |
| 128 | + |
| 129 | + except Exception as e: |
| 130 | + print(f"Warning: Could not build {module_file}: {e}") |
| 131 | + import traceback |
| 132 | + traceback.print_exc() |
| 133 | + |
| 134 | + return built_any |
0 commit comments