pyroxide

11def task(func_or_none=None, *, isolated: bool = False):
12    """
13    Decorator to offload a Python function to the Rust background worker pool.
14
15    The decorated function will be executed on a background OS thread managed
16    by Pyroxide's lock-free task broker. If `isolated=True` is set, the task
17    is executed in a separate OS process, providing full crash safety and GIL-free
18    concurrency for pure Python CPU tasks.
19
20    Args:
21        func_or_none: The Python callable to execute.
22        isolated: Set to True to run the task in an isolated worker process.
23    """
24
25    def decorator(func: Callable[[P], R]) -> Callable[[P], TaskHandle]:
26        @functools.wraps(func)
27        def wrapper(payload: P) -> TaskHandle:
28            import os
29            if os.environ.get("PYROXIDE_WORKER") == "1":
30                return func(payload)
31
32            target_callable = wrapper if isolated else func
33            task_id = submit_task(target_callable, payload, isolated=isolated)
34            return TaskHandle(task_id)
35
36        def batch(payloads: list) -> list[TaskHandle]:
37            from ._pyroxide import submit_batch
38            import os
39            if os.environ.get("PYROXIDE_WORKER") == "1":
40                return [func(p) for p in payloads]
41
42            target_callable = wrapper if isolated else func
43            task_ids = submit_batch(target_callable, payloads, isolated=isolated)
44            return [TaskHandle(tid) for tid in task_ids]
45
46        wrapper.batch = batch
47        return wrapper
48
49    if func_or_none is None:
50        return decorator
51    else:
52        return decorator(func_or_none)

 6class TaskHandle:
 7    def __init__(self, task_id: int) -> None:
 8        self.task_id: int = task_id
 9        self._consumed: bool = False
10
11    def __repr__(self) -> str:
12        return f"<TaskHandle id={self.task_id} status={self.status}>"
13
14    @property
15    def status(self) -> str:
16        """Queries the current status from the Rust Slab."""
17        return get_status(self.task_id)
18
19    def cancel(self) -> bool:
20        """
21        Attempts to cancel the task. Returns True if successfully cancelled,
22        False if the task is already finished or cancelled.
23        """
24        from ._pyroxide import cancel_task
25
26        return cancel_task(self.task_id)
27
28    def wait(
29        self, poll_interval_ms: int = 10, timeout_sec: Optional[float] = None
30    ) -> str:
31        """
32        Blocks the Python runtime until the background Rust worker completes the task.
33        Uses native Rust condvar signal to sleep with 0% CPU usage.
34        """
35        timeout_ms: Optional[int] = (
36            int(timeout_sec * 1000) if timeout_sec is not None else None
37        )
38        current_status: str = wait_status(self.task_id, timeout_ms)
39
40        if timeout_sec is not None and current_status not in ("Completed", "Failed"):
41            raise TimeoutError(f"Task {self.task_id} timed out.")
42
43        return current_status
44
45    def result(self, timeout_sec: Optional[float] = None, consume: bool = True) -> Any:
46        """
47        Blocks until the task is complete, then returns the result.
48        If the task failed, raises the exception encountered.
49
50        Args:
51            timeout_sec: Maximum time in seconds to wait.
52            consume: If True, automatically evicts the task from the Rust Slab once retrieved.
53        """
54        self.wait(timeout_sec=timeout_sec)
55        from ._pyroxide import get_result, free_task
56
57        res = get_result(self.task_id)
58        if consume:
59            free_task(self.task_id)
60            self._consumed = True
61        return res
62
63    async def result_async(
64        self, timeout_sec: Optional[float] = None, consume: bool = True
65    ) -> Any:
66        """
67        Asynchronously awaits the task result, yielding control to the event loop.
68        """
69        import asyncio
70
71        loop = asyncio.get_running_loop()
72        # Offload the blocking wait() call to the loop's default thread pool executor
73        await loop.run_in_executor(None, self.wait, 10, timeout_sec)
74        return self.result(timeout_sec=0, consume=consume)
75
76    def __del__(self) -> None:
77        """
78        Garbage collection destructor.
79        Automatically frees the task memory in the Rust Slab when the Python handle is deleted/dropped.
80        """
81        if getattr(self, "_consumed", False):
82            return
83        try:
84            from ._pyroxide import free_task
85
86            free_task(self.task_id)
87            self._consumed = True
88        except Exception:
89            pass

 7def register_wasm(module_name: str, wasm_bytes: bytes):
 8    """
 9    Registers a pre-compiled WebAssembly module in the global registry.
10    """
11    register_wasm_module(module_name, wasm_bytes)

14def wasm_task(module_name: str, func_name: str = "run", *, isolated: bool = False):
15    """
16    Decorator to submit string or byte payloads to be processed by a registered WASM module.
17    """
18
19    def decorator(func):
20        @functools.wraps(func)
21        def wrapper(payload) -> TaskHandle:
22            task_id = submit_wasm_task(module_name, func_name, payload, isolated=isolated)
23            return TaskHandle(task_id)
24
25        return wrapper
26
27    return decorator

 11def compile_dylib(
 12    name: str, source_code: str, dependencies: Optional[Dict[str, str]] = None
 13) -> str:
 14    """
 15    Compiles Rust source code on-the-fly into a dynamic shared library (.so / .dylib / .dll),
 16    and registers it with the Pyroxide background broker for GIL-free execution.
 17
 18    The compilation is handled automatically by invoking ``cargo build --release`` inside
 19    a temporary directory. The user does not need to install or configure anything beyond
 20    having a working Rust toolchain (``rustc`` + ``cargo``).
 21
 22    Args:
 23        name: Unique name for the dylib. Used to reference it in ``@dylib_task``.
 24        source_code: Raw Rust source code string. Must export two C-compatible symbols:
 25
 26            - ``pyroxide_plugin_run(ptr, len, out_len) -> *mut u8``
 27            - ``pyroxide_plugin_free(ptr, len)``
 28        dependencies: Optional dict of Cargo dependencies, e.g. ``{"serde": "1.0"}``.
 29
 30    Returns:
 31        Absolute path to the compiled shared library file.
 32
 33    Raises:
 34        RuntimeError: If the Cargo compilation fails.
 35        FileNotFoundError: If the compiled library binary is not found after build.
 36
 37    Example:
 38        >>> compile_dylib("my_lib", RUST_SOURCE_CODE)
 39        >>> @dylib_task("my_lib")
 40        ... def process(payload): pass
 41        >>> handle = process("hello")
 42        >>> print(handle.result())
 43    """
 44    temp_dir = tempfile.mkdtemp(prefix=f"pyroxide_dylib_{name}_")
 45    try:
 46        # Run cargo init
 47        subprocess.run(
 48            ["cargo", "init", "--lib", "--name", name],
 49            cwd=temp_dir,
 50            check=True,
 51            stdout=subprocess.DEVNULL,
 52            stderr=subprocess.DEVNULL,
 53        )
 54
 55        cargo_toml_path = os.path.join(temp_dir, "Cargo.toml")
 56        with open(cargo_toml_path, "r") as f:
 57            cargo_content = f.read()
 58
 59        # Force Edition 2021 to prevent newer Rust 2024 edition strict compiler errors
 60        cargo_content = cargo_content.replace('edition = "2024"', 'edition = "2021"')
 61
 62        # Add cdylib configuration
 63        cargo_content += '\n[lib]\ncrate-type = ["cdylib"]\n'
 64
 65        # Add dependencies
 66        if dependencies:
 67            cargo_content += "\n[dependencies]\n"
 68            for dep, ver in dependencies.items():
 69                cargo_content += f'{dep} = "{ver}"\n'
 70
 71        with open(cargo_toml_path, "w") as f:
 72            f.write(cargo_content)
 73
 74        # Write Rust source code to src/lib.rs
 75        lib_rs_path = os.path.join(temp_dir, "src", "lib.rs")
 76        with open(lib_rs_path, "w") as f:
 77            f.write(source_code)
 78
 79        # Run cargo build in release mode
 80        res = subprocess.run(
 81            ["cargo", "build", "--release"],
 82            cwd=temp_dir,
 83            capture_output=True,
 84            text=True,
 85        )
 86        if res.returncode != 0:
 87            raise RuntimeError(f"Cargo build failed:\n{res.stderr}\n{res.stdout}")
 88
 89        # Find compiled library
 90        lib_ext = "dylib" if sys.platform == "darwin" else "so"
 91        if sys.platform == "win32":
 92            lib_ext = "dll"
 93
 94        lib_name = f"lib{name}.{lib_ext}"
 95        if sys.platform == "win32":
 96            lib_name = f"{name}.{lib_ext}"
 97
 98        compiled_path = os.path.join(temp_dir, "target", "release", lib_name)
 99        if not os.path.exists(compiled_path):
100            raise FileNotFoundError(f"Compiled library not found at: {compiled_path}")
101
102        # Register dylib with the Rust core engine
103        register_dylib(name, compiled_path)
104        return compiled_path
105
106    except Exception as e:
107        raise RuntimeError(f"Failed to compile dylib '{name}' via Cargo: {e}") from e

196def dylib_task(dylib_name: str, *, isolated: bool = False):
197    """
198    Decorator that routes task payloads to a registered dynamic shared library (dylib)
199    for GIL-free execution on the background Rust worker pool.
200
201    The dylib must have been previously compiled and registered via ``compile_dylib()``.
202
203    Args:
204        dylib_name: The name of the dylib as registered with ``compile_dylib()``.
205        isolated: Set to True to run in an isolated worker process for crash isolation.
206
207    Example:
208        >>> @dylib_task("my_lib", isolated=True)
209        ... def process(payload: str) -> str:
210        ...     pass  # Execution is handled by the compiled dylib in a separate process
211        >>> handle = process("hello")
212        >>> print(handle.result())
213    """
214
215    def decorator(func: Callable[[Any], Any]) -> Callable[[Any], TaskHandle]:
216        def wrapper(payload: Any) -> TaskHandle:
217            task_id = submit_dylib_task(dylib_name, payload, isolated=isolated)
218            return TaskHandle(task_id)
219
220        def batch(payloads: list) -> list[TaskHandle]:
221            return [wrapper(p) for p in payloads]
222
223        wrapper.batch = batch
224        return wrapper
225
226    return decorator

110def compile_c(name: str, source_code: str) -> str:
111    """
112    Compiles C source code on-the-fly into a dynamic shared library (.so / .dylib / .dll),
113    and registers it with the Pyroxide background broker for GIL-free execution.
114
115    Args:
116        name: Unique name for the library. Used to reference it in @dylib_task.
117        source_code: Raw C source code string. Must export two functions:
118            - ``pyroxide_plugin_run(ptr, len, out_len) -> uint8_t*``
119            - ``pyroxide_plugin_free(ptr, len)``
120    """
121    temp_dir = tempfile.mkdtemp(prefix=f"pyroxide_c_{name}_")
122    try:
123        src_path = os.path.join(temp_dir, f"{name}.c")
124        with open(src_path, "w") as f:
125            f.write(source_code)
126
127        cc = os.environ.get("CC", "clang" if sys.platform == "darwin" else "gcc")
128        lib_ext = "dylib" if sys.platform == "darwin" else "so"
129        if sys.platform == "win32":
130            lib_ext = "dll"
131
132        lib_name = f"lib{name}.{lib_ext}"
133        if sys.platform == "win32":
134            lib_name = f"{name}.{lib_ext}"
135        compiled_path = os.path.join(temp_dir, lib_name)
136
137        cmd = [cc, "-shared", "-o", compiled_path, "-fPIC", src_path]
138        res = subprocess.run(cmd, capture_output=True, text=True)
139        if res.returncode != 0:
140            raise RuntimeError(f"C compilation failed:\n{res.stderr}\n{res.stdout}")
141
142        if not os.path.exists(compiled_path):
143            raise FileNotFoundError(f"Compiled C library not found at: {compiled_path}")
144
145        register_dylib(name, compiled_path)
146        return compiled_path
147
148    except Exception as e:
149        raise RuntimeError(f"Failed to compile C library '{name}': {e}") from e

152def compile_zig(name: str, source_code: str) -> str:
153    """
154    Compiles Zig source code on-the-fly into a dynamic shared library (.so / .dylib / .dll),
155    and registers it with the Pyroxide background broker for GIL-free execution.
156
157    Args:
158        name: Unique name for the library. Used to reference it in @dylib_task.
159        source_code: Raw Zig source code string. Must export two functions:
160            - ``pyroxide_plugin_run(ptr, len, out_len) -> [*]u8``
161            - ``pyroxide_plugin_free(ptr, len)``
162    """
163    temp_dir = tempfile.mkdtemp(prefix=f"pyroxide_zig_{name}_")
164    try:
165        src_path = os.path.join(temp_dir, f"{name}.zig")
166        with open(src_path, "w") as f:
167            f.write(source_code)
168
169        # Compiles dynamic library. Zig build-lib generates output in cwd
170        cmd = ["zig", "build-lib", "-dynamic", "-O", "ReleaseFast", src_path]
171        res = subprocess.run(cmd, cwd=temp_dir, capture_output=True, text=True)
172        if res.returncode != 0:
173            raise RuntimeError(f"Zig compilation failed:\n{res.stderr}\n{res.stdout}")
174
175        lib_ext = "dylib" if sys.platform == "darwin" else "so"
176        if sys.platform == "win32":
177            lib_ext = "dll"
178
179        lib_name = f"lib{name}.{lib_ext}"
180        if sys.platform == "win32":
181            lib_name = f"{name}.{lib_ext}"
182
183        compiled_path = os.path.join(temp_dir, lib_name)
184        if not os.path.exists(compiled_path):
185            raise FileNotFoundError(
186                f"Compiled Zig library not found at: {compiled_path}"
187            )
188
189        register_dylib(name, compiled_path)
190        return compiled_path
191
192    except Exception as e:
193        raise RuntimeError(f"Failed to compile Zig library '{name}': {e}") from e

39def group(handles: Iterable[TaskHandle]) -> TaskGroup:
40    """Wraps multiple task handles into a parallel TaskGroup."""
41    return TaskGroup(handles)

 6class TaskGroup:
 7    """A collection of tasks that run in parallel and can be managed as a unit."""
 8    def __init__(self, handles: Iterable[TaskHandle]):
 9        self.handles = list(handles)
10
11    def __repr__(self) -> str:
12        return f"<TaskGroup handles={self.handles} status={self.status}>"
13
14    @property
15    def status(self) -> str:
16        statuses = [h.status for h in self.handles]
17        if "Failed" in statuses:
18            return "Failed"
19        if "Cancelled" in statuses:
20            return "Cancelled"
21        if all(s == "Completed" for s in statuses):
22            return "Completed"
23        return "Running"
24        
25    def wait(self):
26        """Blocks until all tasks in the group are completed."""
27        for h in self.handles:
28            h.wait()
29            
30    def result(self, consume: bool = True) -> List:
31        """Waits for all tasks and returns their results in order."""
32        return [h.result(consume=consume) for h in self.handles]
33        
34    def cancel(self) -> bool:
35        """Cancels all tasks in the group. Returns True if all were successfully cancelled."""
36        results = [h.cancel() for h in self.handles]
37        return all(results)