Skip to content

Commit 2c9a2fd

Browse files
committed
refactor(core): decouple installers and implement atomic stash-swaps
- Extracted monolithic install logic into and . - Implemented staircase dispatch: FFI → Daemon IPC → subprocess uv → pip. - Added : atomically renames current packages to protect them during uv installs, eliminating the need to rebuild bubbles from scratch. - Updated contract to support atomic plan callback interception.
1 parent eae5bcb commit 2c9a2fd

5 files changed

Lines changed: 2912 additions & 3229 deletions

File tree

Lines changed: 204 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,82 +1,233 @@
11
"""
22
uv_ffi — thin wrapper around the PyO3 native extension.
33
Includes safety stubs for legacy native binaries.
4+
5+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
6+
IPC / FFI CONTRACT OVERVIEW
7+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
8+
All functions below are bound from the PyO3 native extension (_native .so/.pyd).
9+
If the native binary predates a symbol, the getattr() fallback activates instead
10+
of crashing — legacy binaries get _noop / typed stubs.
11+
12+
DATA FLOW (install path):
13+
Python calls run("pip install rich==14.3.3 --target /tmp/bubble_xyz")
14+
15+
▼ [inside Rust / uv]
16+
uv resolves + plans → writes INSTALL_PLAN (Vec<(name, version, action)>)
17+
18+
▼ fires PLAN_READY_CALLBACK (set via set_plan_callback)
19+
20+
▼ Python callback receives plan_entries: list[tuple[str,str,str]]
21+
return True → Rust skips the actual install (Python handled it)
22+
return False → Rust proceeds with normal install
23+
24+
SELF-HEALING RESPONSIBILITY:
25+
When the callback returns True the Rust install() function returns
26+
Changelog::default() immediately. Python then OWNS the operation:
27+
- It must use the SAME target path that was originally passed in run().
28+
- It must clean up stale directories (package dir present, dist-info gone).
29+
- It must not fall back to site-packages if the original target was a bubble.
30+
See: _SelfHealingGuard usage notes below.
31+
32+
PLAN ENTRY TUPLE: (name: str, version: str, action: str)
33+
name — normalised PEP 503 package name e.g. "rich"
34+
version — PEP 440 version string e.g. "14.3.3" (may be "" for remote
35+
entries when uv hasn't resolved the version yet)
36+
action — one of: "cached" | "remote" | "reinstall" | "extraneous"
37+
38+
NOTE: PLAN_HANDLED AtomicBool is REMOVED from the Rust side (was vestigial).
39+
mark_plan_handled() stub is kept here for one-cycle backwards-compat only.
40+
Do NOT use it in new code — return True from the plan callback instead.
41+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
442
"""
543
import sys as _sys
644
import os as _os
745

8-
def _load_native():
9-
_here = _os.path.dirname(__file__)
10-
_native_dir = _os.path.join(_here, '_native')
1146

47+
def _load_native():
1248
import sysconfig as _sc
13-
import glob as _glob
49+
import importlib.util as _ilu
50+
51+
_here = _os.path.dirname(__file__)
1452
_tag = _sc.get_config_var('EXT_SUFFIX') or ''
15-
_versioned = _os.path.join(_native_dir, f'uv_ffi{_tag}')
16-
_abi3 = _os.path.join(_native_dir, 'uv_ffi.abi3.so')
17-
_is_vendored_native = _os.path.exists(_versioned) or _os.path.exists(_abi3)
18-
19-
# Check if this interpreter has its OWN uv_ffi installed — glob for any .so/.pyd
20-
_site_uv_ffi_dir = _os.path.join(_sc.get_path('purelib'), 'uv_ffi')
21-
_site_sos = (
22-
_glob.glob(_os.path.join(_site_uv_ffi_dir, 'uv_ffi*.so')) +
23-
_glob.glob(_os.path.join(_site_uv_ffi_dir, 'uv_ffi*.pyd'))
53+
54+
# abi3 first (everyone except cp37), then cpython-specific (cp37 only)
55+
# same dir first (normal pip install), then _native/ subdir (vendor layout)
56+
_candidates = [
57+
_os.path.join(_here, '_native.abi3.so'),
58+
_os.path.join(_here, f'_native{_tag}'),
59+
_os.path.join(_here, '_native', '_native.abi3.so'),
60+
_os.path.join(_here, '_native', f'_native{_tag}'),
61+
# legacy filenames from before module-name rename
62+
_os.path.join(_here, '_native', 'uv_ffi.abi3.so'),
63+
_os.path.join(_here, '_native', f'uv_ffi{_tag}'),
64+
]
65+
66+
for _so in _candidates:
67+
if _os.path.exists(_so):
68+
_spec = _ilu.spec_from_file_location('uv_ffi._native', _so)
69+
_mod = _ilu.module_from_spec(_spec)
70+
_spec.loader.exec_module(_mod)
71+
return _mod
72+
73+
raise ImportError(
74+
f'uv_ffi: no .so found for {_sys.executable} '
75+
f'(tag={_tag}, site={_sc.get_path("purelib")})'
2476
)
25-
_sys.stderr.flush()
26-
27-
if _site_sos:
28-
_site_so = _site_sos[0]
29-
import importlib.util as _ilu
30-
_spec = _ilu.spec_from_file_location("uv_ffi", _site_so)
31-
_mod = _ilu.module_from_spec(_spec)
32-
_spec.loader.exec_module(_mod)
33-
_sys.stderr.flush()
34-
return _mod
35-
36-
if _is_vendored_native:
37-
_so = _versioned if _os.path.exists(_versioned) else _abi3
38-
import importlib.util as _ilu
39-
_spec = _ilu.spec_from_file_location("uv_ffi._native", _so)
40-
_mod = _ilu.module_from_spec(_spec)
41-
_spec.loader.exec_module(_mod)
42-
_sys.stderr.flush()
43-
return _mod
44-
45-
raise ImportError(f'uv_ffi: no .so found for {_sys.executable} (tag={_tag}, site={_sc.get_path("purelib")})')
4677

4778
# Load the native binary
4879
_native = _load_native()
4980
_loaded_so_path = getattr(_native, '__file__', 'unknown')
5081

51-
# ── SAFETY STUBS ─────────────────────────────────────────────────────────────
52-
# We use getattr(..., fallback) to ensure that if a legacy .so is loaded,
53-
# the import doesn't crash with an AttributeError.
54-
# The functions will exist as no-ops instead of blowing up the process.
82+
# ── SAFETY STUBS ──────────────────────────────────────────────────────────────
83+
# getattr(obj, name, fallback) means: if the loaded .so is older and doesn't
84+
# export the symbol yet, use the fallback instead of raising AttributeError.
5585
# ─────────────────────────────────────────────────────────────────────────────
5686
_noop = lambda *a, **kw: None
5787

58-
run = getattr(_native, 'run', lambda cmd: (1, "Native 'run' missing", ""))
88+
89+
# ── CORE EXECUTION ────────────────────────────────────────────────────────────
90+
91+
# run(cmd: str) -> tuple[int, Changelog | None, str] [from Rust: (i32, Option<Changelog>, String)]
92+
# Exposed through the run() function below which normalises to 4-tuple.
93+
# Do NOT call _native.run() directly — always use run() / run_capture().
94+
95+
# ── CACHE MANAGEMENT ─────────────────────────────────────────────────────────
96+
97+
# get_site_packages_cache() -> object | None
98+
# Returns the cached Arc<SitePackages> handle (opaque Python object).
99+
# Returns None if the cache has not been primed.
59100
get_site_packages_cache = getattr(_native, 'get_site_packages_cache', _noop)
101+
102+
# invalidate_site_packages_cache() -> None
103+
# Drops the cached Arc<SitePackages>, forcing a fresh scan on next access.
60104
invalidate_site_packages_cache = getattr(_native, 'invalidate_site_packages_cache', _noop)
105+
106+
# patch_site_packages_cache(dist_name: str, version: str, location: str) -> None
107+
# Surgically insert one dist into the live cache without a full rescan.
108+
# dist_name : normalised PEP 503 name ("rich", "numpy")
109+
# version : PEP 440 string ("14.3.3")
110+
# location : absolute path to the dist-info directory
61111
patch_site_packages_cache = getattr(_native, 'patch_site_packages_cache', _noop)
112+
113+
# clear_registry_cache() -> None
114+
# Wipes the uv registry / index response cache in-process.
115+
# Call before any install where you suspect stale PyPI metadata.
116+
# Added in uv-ffi post6 — older binaries get _noop.
62117
clear_registry_cache = getattr(_native, 'clear_registry_cache', _noop)
63118

119+
# ── BUBBLE CACHE ─────────────────────────────────────────────────────────────
120+
121+
# evict_bubble_cache() -> None
122+
# Clears ALL bubble site-packages cache entries.
123+
# Use when a bubble is destroyed or its contents change externally.
124+
evict_bubble_cache = getattr(_native, 'evict_bubble_cache', _noop)
125+
126+
# evict_packages_from_bubble_cache(bubble_id: str, *package_names: str) -> None
127+
# Evict specific packages from a single bubble's cache entry.
128+
# bubble_id : the bubble identifier string used at creation time
129+
# package_names : one or more normalised PEP 503 package names
130+
evict_packages_from_bubble_cache = getattr(_native, 'evict_packages_from_bubble_cache', _noop)
131+
132+
# patch_bubble_site_packages_cache(bubble_id: str, dist_name: str,
133+
# version: str, location: str) -> None
134+
# Equivalent of patch_site_packages_cache but for a named bubble.
135+
# bubble_id : the bubble identifier string
136+
# dist_name : normalised PEP 503 name
137+
# version : PEP 440 string
138+
# location : absolute path to the dist-info directory inside the bubble
139+
patch_bubble_site_packages_cache = getattr(_native, 'patch_bubble_site_packages_cache', _noop)
140+
141+
# ── INSTALL PLAN IPC ─────────────────────────────────────────────────────────
142+
143+
# get_install_plan() -> list[tuple[str, str, str]]
144+
# Returns the last plan written by Rust during the current or most-recent
145+
# install() call. Each entry is (name, version, action).
146+
# action is one of: "cached" | "remote" | "reinstall" | "extraneous"
147+
# version may be "" for "remote" entries (version not yet resolved).
148+
# The list is CLEARED at the start of each resolve() and install() call,
149+
# and again at the start of execute_plan() sub-steps.
150+
#
151+
# Thread safety: backed by std::sync::Mutex — safe to call from any thread,
152+
# but reads are not atomic with respect to concurrent Rust writes.
153+
# Best practice: read only inside the plan callback, not after returning.
154+
#
155+
# Fallback on legacy .so: returns []
156+
get_install_plan = getattr(_native, 'get_install_plan', lambda: [])
157+
158+
# set_plan_callback(cb: Callable[[list[tuple[str, str, str]]], bool]) -> None
159+
# Register a Python callable that Rust fires when the install plan is ready.
160+
#
161+
# cb signature:
162+
# def my_callback(entries: list[tuple[str, str, str]]) -> bool:
163+
# ...
164+
# return handled # True → Rust skips install; False → Rust proceeds
165+
#
166+
# entries — same format as get_install_plan() return value (snapshot at
167+
# callback time, same list that was just written to INSTALL_PLAN)
168+
#
169+
# ⚠️ CRITICAL — TARGET TRACKING:
170+
# When the callback returns True, Rust exits install() immediately.
171+
# The callback is responsible for performing any needed operation on
172+
# the SAME target that was passed to run(). Do NOT default to
173+
# site-packages — always carry the original target path through
174+
# your callback closure.
175+
#
176+
# ⚠️ CRITICAL — STALE DIR CLEANUP:
177+
# Before installing into a target dir, check for the case where the
178+
# package directory exists but its dist-info is gone (uv partial-write
179+
# or interrupted uninstall). If found: remove the package dir first,
180+
# then install fresh. uv may warn "files may have been left behind" —
181+
# trust that warning and act on it.
182+
#
183+
# Only one callback can be registered at a time (last write wins).
184+
# Fallback on legacy .so: _noop (callback never fires, plan never handled).
185+
set_plan_callback = getattr(_native, 'set_plan_callback', _noop)
186+
187+
# mark_plan_handled() -> None
188+
# ⚠️ DEPRECATED — do NOT use in new code.
189+
# The PLAN_HANDLED AtomicBool has been removed from the Rust side.
190+
# This stub exists only to avoid ImportError on one-cycle-old callers.
191+
# Signal "handled" by returning True from your set_plan_callback callback.
192+
mark_plan_handled = getattr(_native, 'mark_plan_handled', _noop)
193+
194+
195+
# ── CORE run() WRAPPER ────────────────────────────────────────────────────────
196+
64197
def run(cmd: str) -> tuple:
65198
"""
66-
Core execution wrapper.
67-
If the native 'run' was missing, this will call the stub return value.
199+
Execute a uv command via the Rust FFI layer.
200+
201+
Args:
202+
cmd: Shell-style uv argument string, WITHOUT the leading "uv".
203+
Examples:
204+
"pip install rich==14.3.3"
205+
"pip install rich==15.0.0 --target /tmp/bubble_xyz"
206+
"pip uninstall rich --yes"
207+
208+
Returns:
209+
4-tuple: (returncode: int, changelog: object | None,
210+
stderr: str, extra: str)
211+
returncode == 0 → success
212+
returncode != 0 → failure; check stderr for details
213+
214+
Notes:
215+
- If set_plan_callback was set and the callback returns True, the
216+
install is considered handled by Python. returncode will be 0
217+
and changelog will be a default (empty) Changelog.
218+
- Always call run() through this wrapper, not _native.run() directly,
219+
so that the 3-tuple → 4-tuple normalisation is applied consistently.
68220
"""
69-
# If 'run' is the stub, it returns a tuple; otherwise call the native function.
70221
if hasattr(_native, 'run'):
71222
result = _native.run(cmd)
72223
if len(result) == 3:
73224
return (result[0], result[1], result[2], '')
74225
return result
75226
else:
76-
# Fallback return for the stub to prevent crashing the caller
77-
return (1, "Native 'run' method not found in loaded .so", "")
227+
return (1, None, "Native 'run' method not found in loaded .so", '')
228+
78229

79-
run_capture = run
230+
run_capture = run # alias kept for callers using the older name
80231

81232
try:
82233
from importlib.metadata import version as _meta_version
@@ -85,10 +236,17 @@ def run(cmd: str) -> tuple:
85236
__version__ = "unknown"
86237

87238
__all__ = [
88-
"run", "run_capture",
239+
"run",
240+
"run_capture",
89241
"get_site_packages_cache",
90242
"invalidate_site_packages_cache",
91243
"patch_site_packages_cache",
92244
"clear_registry_cache",
245+
"evict_bubble_cache",
246+
"evict_packages_from_bubble_cache",
247+
"patch_bubble_site_packages_cache",
248+
"get_install_plan",
249+
"set_plan_callback",
250+
"mark_plan_handled", # deprecated stub — see docstring above
93251
"__version__",
94252
]

0 commit comments

Comments
 (0)