okami.mom/albedo
git clone https://git.nkp.pet/okami.mom/albedo
74cf623
main
1"""albedo's nix: the program behind the `nix` shim (albedo_shims.py) that 2the model's jobs find first on PATH. 3 4`nix` in a jj workspace that has no .git hands nix a plain path, and nix 5copies that whole directory into the store, build and target dirs included, 6on every evaluation. The shim names the workspace's commit in the 7repository's git store instead (git+file://...?rev=...), so nix reads tracked 8files only, and points the lock file nix may write back at the workspace 9(--output-lock-file), where nix would have put it in a git checkout. 10 11`nix develop [installable] -c program ...` runs the program in a dev shell 12captured once, as `nix develop` would build it (the dev shell script and its 13shellHook), and kept as the difference it makes to the environment. It is 14cached under $ALBEDO_HOME/cache/nix-develop by the content of the files dev 15shells are made of (SHELL_FILES), behind a profile nix's garbage collector 16keeps, so an unchanged flake costs no nix call in any checkout of the 17repository. 18 19Commands are read with nix's own flag table (`nix __dump-cli`, cached per nix 20binary); everything the shim has no reason to change goes to nix unchanged. 21ALBEDO_PLAIN_NIX=1 runs nix as is. 22""" 23 24from __future__import annotations 25 26import fcntl 27import fnmatch 28import hashlib 29import json 30import os 31import platform 32import re 33import shutil 34import subprocess 35import sys 36import tempfile 37import time 38from dataclasses import dataclass 39from pathlib import Path 40from typing import Any ,NoReturn 41from urllib .parse import quote 42 43FEATURES = ["--extra-experimental-features" ,"nix-command flakes" ] 44# What a dev shell is made of; an edit anywhere else keeps the cached shell. 45SHELL_FILES = ("*.nix" ,"flake.lock" ,"rust-toolchain" ,"rust-toolchain.toml" ) 46# Never walked for SHELL_FILES in a directory outside version control. 47UNWALKED = {"node_modules" ,"target" ,"build" ,"_build" ,"dist" ,"result" } 48CACHE_LIMIT = 32 # cached shells; the least recently used goes first 49REMOTE_TTL = 86400 # seconds a shell from a flake elsewhere (github:...) is kept 50# the capture a record's delta describes; an older one missed exports and is 51# rebuilt rather than replayed 52FORMAT = "exports" 53STEP_TIMEOUT = 600 # seconds for one jj or git call, or a shellHook 54# Set by sourcing the shell script, meaningless outside the bash that did it. 55VOLATILE = {"NIX_BUILD_TOP" ,"TMP" ,"TMPDIR" ,"TEMP" ,"TEMPDIR" } 56VOLATILE |= {"PWD" ,"OLDPWD" ,"SHLVL" ,"_" } 57# the nix compiler wrapper's inputs, which no build cache key covers (albedo_shims) 58WRAPPER = ("NIX_CFLAGS" ,"NIX_LDFLAGS" ,"NIX_CC" ,"NIX_HARDENING_ENABLE" ) 59# develop flags a cached shell honours: they change what nix prints, not the shell 60QUIET = {"verbose" ,"quiet" ,"print-build-logs" ,"log-format" ,"show-trace" } 61QUIET |= {"debug" ,"accept-flake-config" ,"no-write-lock-file" ,"no-update-lock-file" } 62# flags that write into the repository the ref names, not into the workspace 63WRITING = {"commit-lock-file" } 64# nix's own spellings of commands the flag table lists under another name, 65# by the command they are spelled under 66ALIASES = { 67 (): {"shell" : ["env" ,"shell" ],"dev-shell" : ["develop" ]}, 68 ("profile" ,): {"install" : ["add" ]}, 69} 70# argument labels (nix __dump-cli) of one word naming a flake or its installable; 71# "installables" is every word 72FLAKE_ARGS = {"installable" ,"flake-url" ,"package" ,"dependency" } 73VARIADIC = {"args" ,"strings" ,"paths" ,"elements" ,"hashes" ,"inputs" } 74# flags that make installables something other than flakes 75NOT_FLAKES = {"expr" ,"file" } 76# commands that run the formatter of the flake at `.`, and the nix verb for that 77FORMATTERS = { 78 ("fmt" ,):"run" , 79 ("formatter" ,"run" ):"run" , 80 ("formatter" ,"build" ):"build" , 81} 82NOTE = "albedo:" 83REFUSED = ( 84"nix would copy this jj workspace (it has no .git) into the store, build " 85"output included. Name the flake by its commit (git+file://<repo>?rev=<commit>) " 86"or set ALBEDO_PLAIN_NIX=1 to copy anyway." 87) 88 89# what a dev shell does to an environment: "set", "prepend", and "unset" 90Change = dict [str ,Any ] 91 92 93class NixError (RuntimeError ): 94"""jj, git, or nix could not prepare the shell; the message has their output.""" 95 96 97class Unreadable (Exception ): 98"""A flag or command nix's own table does not have: nix refuses it 99itself, or it is an alias the shim does not know.""" 100 101 102@dataclass( frozen = True ) 103class Flake : 104"""A local flake and the checkout it lives in.""" 105 106directory :Path # holds flake.nix 107root :Path # the jj workspace, the git work tree, or the directory itself 108vcs :str # "jj", "git", or "" outside version control 109 110@property 111def copied (self )-> bool : 112"""Whether nix, handed this directory, copies all of it: a jj 113workspace with no .git beside .jj is a plain path to nix.""" 114return self .vcs == "jj" and not (self .root / ".git" ).exists () 115 116 117def _above (start :Path ,name :str )-> Path | None : 118for directory in (start ,* start .parents ): 119if (directory / name ).exists (): 120return directory 121return None 122 123 124def locate (start :Path )-> Flake | None : 125"""The flake at or above `start`. jj wins over git at the same root, as 126the picker reads a colocated repository.""" 127directory = _above (start ,"flake.nix" ) 128if directory is None : 129return None 130jj ,git = _above (directory ,".jj" ),_above (directory ,".git" ) 131if jj is not None and (git is None or git == jj or git in jj .parents ): 132return Flake (directory ,jj ,"jj" ) 133if git is not None : 134return Flake (directory ,git ,"git" ) 135return Flake (directory ,directory ,"" ) 136 137 138def system ()-> str : 139"""This host's nix system, as nix names it.""" 140machine = {"arm64" :"aarch64" ,"amd64" :"x86_64" }.get (platform .machine ().lower ()) 141kernel = "darwin" if sys .platform == "darwin" else "linux" 142return f" { machine or platform . machine () } - { kernel } " 143 144 145def _local (base :str )-> bool : 146"""Whether an installable's flake part names a local directory.""" 147return base in ("" ,"." ,".." )or base .startswith (("./" ,"../" ,"/" ,"path:" )) 148 149 150def _step (argv :list [str ],cwd :Path | None ,what :str )-> str : 151"""One program the shim needs the output of, or a NixError with it.""" 152try : 153done = subprocess .run ( 154argv ,cwd = cwd ,capture_output = True ,text = True ,timeout = STEP_TIMEOUT 155 ) 156except (OSError ,subprocess .TimeoutExpired )as error : 157raise NixError (f" { what } failed: { error } " )from error 158if done .returncode != 0 : 159raise NixError ( 160f" { what } failed (exit { done . returncode } ):\n { done . stderr [ - 4000 :] } " 161 ) 162return done .stdout 163 164 165def _home ()-> Path : 166return Path (os .environ .get ("ALBEDO_HOME" )or os .path .expanduser ("~/.albedo" )) 167 168 169# ---- reading a command line ---- 170 171 172def flag_table (nix :str )-> dict [str ,Any ]: 173"""nix's commands and the arity of each flag, from `nix __dump-cli`, 174kept per nix binary.""" 175binary = os .path .realpath (nix ) 176cache = _home ()/ "cache" / "nix-cli" 177entry = cache / (hashlib .sha256 (binary .encode ()).hexdigest ()[:32 ]+ ".json" ) 178try : 179return json .loads (entry .read_text ()) 180except (OSError ,ValueError ): 181pass 182dump = json .loads (_step ([binary ,"__dump-cli" ],None ,"nix __dump-cli" ))["args" ] 183 184def node (command :dict [str ,Any ])-> dict [str ,Any ]: 185flags = command .get ("flags" , {}) 186return { 187"flags" : {name :flag .get ("arity" )or 0 for name ,flag in flags .items ()}, 188"short" : { 189flag ["shortName" ]:name 190for name ,flag in flags .items () 191if flag .get ("shortName" ) 192 }, 193"args" : [arg .get ("label" ,"" )for arg in command .get ("args" , [])], 194"commands" : { 195name :node (sub )for name ,sub in command .get ("commands" , {}).items () 196 }, 197 } 198 199table = node (dump ) 200cache .mkdir (parents = True ,exist_ok = True ) 201temporary = cache / f". { entry . name } . { os . getpid () } " 202temporary .write_text (json .dumps (table )) 203temporary .replace (entry ) 204return table 205 206 207@dataclass 208class Command : 209"""One nix command line, read with the flag table.""" 210 211words :list [str ] 212path :list [str ]# the subcommand, e.g. ["flake", "check"] 213named :list [int ]# indices of the words that spell the subcommand 214positional :list [int ]# indices of the words that are not flags or values 215flags :dict [str ,list [int ]]# flag name -> index of each use 216args :list [str ]# the subcommand's argument labels 217program :list [str ]# what follows --command / -c (develop, shell) 218stop :int # index of `--`, or of the end 219 220@property 221def passed (self )-> list [str ]: 222"""What follows `--`.""" 223return self .words [self .stop + 1 :] 224 225@property 226def end (self )-> int : 227"""Where flags for the subcommand can go: just past its words.""" 228return self .named [- 1 ]+ 1 if self .named else 0 229 230 231def read (words :list [str ],table :dict [str ,Any ])-> Command : 232"""The command `words` (nix's arguments) spell, or Unreadable.""" 233node ,path ,named = table , [], [] 234positional :list [int ]= [] 235flags :dict [str ,list [int ]]= {} 236program :list [str ]= [] 237index = 0 238while index < len (words ): 239word = words [index ] 240if word == "--" : 241break 242if word .startswith ("--" ): 243name = word [2 :] 244elif word .startswith ("-" )and len (word )== 2 : 245name = node ["short" ].get (word [1 ])or table ["short" ].get (word [1 ],"" ) 246elif word .startswith ("-" )and word != "-" : 247raise Unreadable (word ) 248else : 249names = ALIASES .get (tuple (path ), {}).get (word , [word ]) 250if not positional and names [0 ]in node ["commands" ]: 251for name in names : 252node = node ["commands" ][name ] 253path .extend (names ) 254named .append (index ) 255elif node ["commands" ]: 256raise Unreadable (word )# a command nix spells some other way 257else : 258positional .append (index ) 259index += 1 260continue 261if name == "command" and path in (["develop" ], ["env" ,"shell" ]): 262program = words [index + 1 :] 263flags .setdefault (name , []).append (index ) 264index = len (words ) 265break 266arity = node ["flags" ].get (name ,table ["flags" ].get (name )) 267if arity is None : 268raise Unreadable (word ) 269flags .setdefault (name , []).append (index ) 270index += 1 + arity 271return Command (words ,path ,named ,positional ,flags ,node ["args" ],program ,index ) 272 273 274def _flake_words (command :Command )-> list [int ]: 275"""Indices of the words that name a flake (or an installable of one), 276following the subcommand's argument labels.""" 277if NOT_FLAKES & command .flags .keys (): 278return []# installables are attribute paths or store paths then 279found ,left = [],list (command .positional ) 280for label in command .args : 281if not left : 282break 283if label == "installables" : 284found ,left = found + left , [] 285elif label in FLAKE_ARGS : 286found .append (left .pop (0 )) 287elif label in VARIADIC : 288left = [] 289else : 290left .pop (0 ) 291return found 292 293 294# ---- refs that never copy a jj workspace ---- 295 296 297def flake_ref (flake :Flake )-> str : 298"""The ref nix reads `flake` through without copying the working copy.""" 299if flake .vcs == "git" : 300return str (flake .directory ) 301if flake .vcs != "jj" : 302return f"path: { flake . directory } " 303# jj log snapshots the working copy, so the commit holds every edit 304log = _step ( 305 ["jj" ,"log" ,"-r" ,"@" ,"--no-graph" ,"-T" ,'"commit=" ++ commit_id ++ "\\n"' ], 306flake .root , 307"jj log" , 308 ) 309commit = re .search (r"commit=([0-9a-f]{40})" ,log ) 310if commit is None : 311raise NixError (f"jj log named no commit:\n { log [ - 2000 :] } " ) 312lines = _step (["jj" ,"git" ,"root" ],flake .root ,"jj git root" ).splitlines () 313store = Path ([line for line in lines if line .strip ()][- 1 ].strip ()) 314# a colocated .git is read through its work tree, a bare store as itself 315repository = store .parent if store .name == ".git" else store 316ref = f"git+file:// { quote ( str ( repository )) } ?rev= { commit . group ( 1 ) } " 317relative = flake .directory .relative_to (flake .root ) 318return ref if relative == Path ("." )else f" { ref } &dir= { quote ( relative . as_posix ()) } " 319 320 321def _copied (word :str ,cwd :Path )-> Flake | None : 322"""The flake `word` names when it is a local one nix would copy.""" 323base = word .partition ("#" )[0 ] 324if not _local (base ): 325return None 326flake = locate ((cwd / base .removeprefix ("path:" )).resolve ()) 327return flake if flake is not None and flake .copied else None 328 329 330def _formatter ( 331command :Command ,flake :Flake ,verb :str 332)-> tuple [list [str ],dict [str ,str ]]: 333"""`nix fmt` and `nix formatter run|build` as the `nix run|build` of the 334workspace flake's formatter, run from the flake's root as nix fmt does.""" 335words = command .words 336skipped = {* command .named ,* command .positional } 337keep = [word for i ,word in enumerate (words [:command .stop ])if i not in skipped ] 338passed = [words [i ]for i in command .positional ]+ command .passed 339lock = ["--output-lock-file" ,str (flake .directory / "flake.lock" )] 340rewritten = [ 341verb , 342* lock , 343f" { flake_ref ( flake ) } #formatter. { system () } " , 344* keep , 345* (["--" ,* passed ]if passed and verb == "run" else []), 346 ] 347return rewritten , {"PRJ_ROOT" :str (flake .directory )} 348 349 350def rewritten (command :Command ,cwd :Path )-> tuple [list [str ],dict [str ,str ]]| None : 351"""`command` with every local flake of a jj workspace named by its 352commit, and the environment changes that go with it; None when no word 353names such a flake. 354 355nix writes a lock file it had to create or update into the repository 356the ref names, the main checkout, so the workspace's own goes to the 357workspace (--output-lock-file); that flag covers every flake the command 358locks, so a command that locks another flake too writes none.""" 359verb = FORMATTERS .get (tuple (command .path )) 360if verb is not None : 361flake = _copied ("." ,cwd ) 362return None if flake is None else _formatter (command ,flake ,verb ) 363words = list (command .words ) 364targets = _flake_words (command ) 365if command .path == ["flake" ,"update" ]: 366targets = [index + 1 for index in command .flags .get ("flake" , [])] 367takes = command .path == ["flake" ,"update" ]or any ( 368label == "installables" or label in FLAKE_ARGS for label in command .args 369 ) 370# nix reads `.` when no flake is named, except for nix eval 371implicit = ( 372_copied ("." ,cwd ) 373if takes 374and not targets 375and (command .path == ["flake" ,"update" ]or not command .positional ) 376and command .path != ["eval" ] 377and not (NOT_FLAKES | {"stdin" ,"all" })& command .flags .keys () 378else None 379 ) 380rewrites :dict [int ,Flake ]= {} 381for index in [* targets ,* (i + 1 for i in command .flags .get ("inputs-from" , []))]: 382flake = _copied (words [index ],cwd ) 383if flake is not None : 384rewrites [index ]= flake 385if not rewrites and implicit is None : 386return None 387if WRITING & command .flags .keys (): 388raise NixError ( 389f"--commit-lock-file would commit in the main checkout; { REFUSED } " 390 ) 391for index ,flake in rewrites .items (): 392_ ,hash_ ,attribute = words [index ].partition ("#" ) 393words [index ]= flake_ref (flake )+ hash_ + attribute 394locked = {rewrites [index ].directory for index in targets if index in rewrites } 395others = [index for index in targets if index not in rewrites ] 396if implicit is not None : 397ref = flake_ref (implicit ) 398flag = ["--flake" ,ref ]if command .path == ["flake" ,"update" ]else [ref ] 399words [command .end :command .end ]= flag 400locked = {implicit .directory } 401if not {"output-lock-file" ,"no-write-lock-file" }& command .flags .keys (): 402if len (locked )== 1 and not others : 403lock = ["--output-lock-file" ,str (locked .pop ()/ "flake.lock" )] 404else : 405lock = ["--no-write-lock-file" ] 406words [command .end :command .end ]= lock 407return words , {} 408 409 410# ---- dev shells ---- 411 412 413def _shell_files (flake :Flake )-> list [str ]: 414"""The checkout's SHELL_FILES, relative to its root: what jj or git 415tracks, or a walk outside version control.""" 416if flake .vcs == "jj" : 417fileset = " | " .join (f'glob:"**/ { pattern } "' for pattern in SHELL_FILES ) 418output = _step ( 419 ["jj" ,"file" ,"list" ,"--ignore-working-copy" ,fileset ], 420flake .root , 421"jj file list" , 422 ) 423return output .splitlines () 424if flake .vcs == "git" : 425specs = [f":(glob)**/ { pattern } " for pattern in SHELL_FILES ] 426output = _step ( 427 ["git" ,"ls-files" ,"-z" ,"--" ,* specs ],flake .root ,"git ls-files" 428 ) 429return [path for path in output .split ("\0" )if path ] 430found = [] 431for directory ,names ,entries in os .walk (flake .root ): 432names [:]= [n for n in names if not n .startswith ("." )and n not in UNWALKED ] 433found += [ 434os .path .relpath (os .path .join (directory ,entry ),flake .root ) 435for entry in entries 436if any (fnmatch .fnmatch (entry ,pattern )for pattern in SHELL_FILES ) 437 ] 438return found 439 440 441def _key (flake :Flake ,attribute :str )-> str : 442"""Names the shell by what makes it: the system, the attribute, where the 443flake sits in its checkout, and the content of every shell file.""" 444digest = hashlib .sha256 () 445relative = flake .directory .relative_to (flake .root ).as_posix () 446for part in (system (),attribute ,relative ): 447digest .update (part .encode ()+ b"\0" ) 448for path in sorted (_shell_files (flake )): 449try : 450content = (flake .root / path ).read_bytes () 451except OSError : 452continue # tracked, but deleted in the working copy 453digest .update (path .encode ()+ b"\0" + hashlib .sha256 (content ).digest ()) 454return digest .hexdigest ()[:32 ] 455 456 457# one `export NAME` of a dev shell script, and one of its `nix_saved_NAME` 458EXPORT = re .compile (r"^export ([A-Za-z_][A-Za-z0-9_]*)(?:[ \t]*[=;]|[ \t]*$)" ,re .M ) 459SAVED = re .compile (r"^nix_saved_([A-Za-z_][A-Za-z0-9_]*)=" ,re .M ) 460 461 462def exported (script :str )-> frozenset [str ]: 463"""The variables a dev shell gives every program in it: what its script 464exports at the top level (the function bodies do not run), minus the 465search paths it composes from the caller's own, saved as 466`nix_saved_<name>`. `nix develop` exports these whatever the caller's 467environment held, so a cached shell carries them all.""" 468return frozenset (EXPORT .findall (script ))- frozenset (SAVED .findall (script )) 469 470 471def delta ( 472before :dict [str ,str ], 473after :dict [str ,str ], 474provided :frozenset [str ]= frozenset (), 475)-> Change : 476"""What a dev shell did to an environment: variables it set, prefixes it 477put before a list it kept (PATH), and variables it unset. `provided` is 478what the shell gives every program (exported()); those stay in even where 479the capture's own environment already held the same value, which is how a 480cached shell came to drop mkShell's env.* values.""" 481changed :dict [str ,str ]= {} 482prepended :dict [str ,str ]= {} 483for name ,value in after .items (): 484old = before .get (name ) 485if name in VOLATILE or (old == value and name not in provided ): 486continue 487if old and value .endswith (":" + old ): 488prepended [name ]= value [:- len (old )- 1 ] 489else : 490changed [name ]= value 491unset = [ 492n for n in before if n not in after and n not in VOLATILE and n .isidentifier () 493 ] 494return {"set" :changed ,"prepend" :prepended ,"unset" :unset } 495 496 497def applied (change :Change ,base :dict [str ,str ])-> dict [str ,str ]: 498"""`base` as the dev shell leaves it.""" 499unset = set (change .get ("unset" , ())) 500environment = {name :value for name ,value in base .items ()if name not in unset } 501environment .update (change .get ("set" , {})) 502for name ,prefix in change .get ("prepend" , {}).items (): 503environment [name ]= ( 504f" { prefix } : { environment [ name ] } " if environment .get (name )else prefix 505 ) 506return environment 507 508 509def capture (nix :str ,arguments :list [str ],profile :Path ,cwd :Path | None )-> Change : 510"""Build the dev shell, source it with its shellHook, and keep what it 511changed. The profile is a garbage-collector root for the shell. nix's 512progress and the hook's output go to stderr.""" 513profile .parent .mkdir (parents = True ,exist_ok = True ) 514with tempfile .TemporaryDirectory (prefix = "albedo-nix-" )as scratch : 515script = Path (scratch ,"rc" ) 516with script .open ("w" )as out : 517built = subprocess .run ( 518 [ 519nix , 520"print-dev-env" , 521* FEATURES , 522"--profile" , 523str (profile ), 524* arguments , 525 ], 526cwd = cwd , 527stdout = out , 528 ) 529if built .returncode != 0 : 530raise NixError (f"nix print-dev-env failed (exit { built . returncode } )" ) 531provided = exported (script .read_text ()) 532variables = json .loads (Path (os .path .realpath (profile )).read_text ()).get ( 533"variables" , {} 534 ) 535bash = variables .get ("shell" , {}).get ("value" ,"" ) 536if os .path .basename (bash )!= "bash" : 537bash = shutil .which ("bash" )or "bash" # the script needs bash 4 or later 538sourced = subprocess .run ( 539 [ 540bash , 541"--noprofile" , 542"--norc" , 543"-c" , 544'exec 3>&1 1>&2; . "$1"; rm -rf "$NIX_BUILD_TOP"; env -0 >&3' , 545"bash" , 546str (script ), 547 ], 548cwd = cwd , 549stdout = subprocess .PIPE , 550timeout = STEP_TIMEOUT , 551 ) 552if sourced .returncode != 0 : 553raise NixError ( 554f"the dev shell's shellHook failed (exit { sourced . returncode } )" 555 ) 556after = dict ( 557item .split ("=" ,1 ) 558for item in sourced .stdout .decode ("utf-8" ,"surrogateescape" ).split ("\0" ) 559if "=" in item 560 ) 561change = delta (dict (os .environ ),after ,provided ) 562wrapper = sorted ((n ,v )for n ,v in after .items ()if n .startswith (WRAPPER )) 563if wrapper :# the shell's compiler-wrapper flags, for the cargo shim 564digest = hashlib .sha256 (json .dumps (wrapper ).encode ()).hexdigest ()[:16 ] 565change ["set" ]["ALBEDO_NIX_ENV" ]= digest 566return change 567 568 569def _cache ()-> Path : 570return _home ()/ "cache" / "nix-develop" 571 572 573def _cached (key :str ,ttl :float | None )-> dict [str ,Any ]| None : 574entry = _cache ()/ f" { key } .json" 575try : 576record = json .loads (entry .read_text ()) 577except (OSError ,ValueError ): 578return None 579if record .get ("format" )!= FORMAT : 580return None 581profile = _cache ()/ record .get ("profile" ,key ) 582stale = ttl is not None and time .time ()- float (record .get ("created" ,0 ))> ttl 583if stale or not os .path .exists (profile ):# collected since, or never built 584return None 585os .utime (entry ) 586return record 587 588 589def _keep (key :str ,record :dict [str ,Any ])-> None : 590"""Store a shell, then drop the least recently used past CACHE_LIMIT.""" 591cache = _cache () 592temporary = cache / f". { key } . { os . getpid () } " 593temporary .write_text (json .dumps ({** record ,"created" :time .time ()})) 594temporary .replace (cache / f" { key } .json" ) 595entries = sorted (cache .glob ("*.json" ),key = lambda path :path .stat ().st_mtime ) 596for old in entries [:- CACHE_LIMIT ]: 597profile = cache / old .stem 598for path in [ 599old , 600profile , 601profile .with_suffix (".lock" ), 602* cache .glob (f" { old . stem } -*-link" ), 603 ]: 604path .unlink (missing_ok = True ) 605 606 607def dev_shell (nix :str ,installable :str ,cwd :Path )-> Change | None : 608"""The cached dev shell for `installable`, built first if it is not 609cached; None for a local installable with no flake.nix. Two shims asking 610for the same shell at once share one build.""" 611base ,hash_ ,attribute = installable .partition ("#" ) 612if _local (base ): 613flake = locate ((cwd / base .removeprefix ("path:" )).resolve ()) 614if flake is None : 615return None 616key ,ttl ,where = _key (flake ,attribute ),None ,flake .directory 617else :# nixpkgs#hello, github:owner/repo: nothing to copy 618flake ,ttl ,where = None ,REMOTE_TTL ,None 619key = hashlib .sha256 (f" { system () } \0 { installable } " .encode ()).hexdigest ()[:32 ] 620record = _cached (key ,ttl ) 621if record is not None : 622return record ["delta" ] 623_cache ().mkdir (parents = True ,exist_ok = True ) 624with (_cache ()/ f" { key } .lock" ).open ("w" )as lock : 625fcntl .flock (lock ,fcntl .LOCK_EX ) 626record = _cached (key ,ttl )# built while this one waited 627if record is None : 628arguments = [installable ] 629if flake is not None : 630ref = flake_ref (flake ) 631arguments = [ref + hash_ + attribute ] 632if flake .copied : 633arguments += [ 634"--output-lock-file" , 635str (flake .directory / "flake.lock" ), 636 ] 637f" { NOTE } building the dev shell for { where or installable } ; " 639"later runs reuse it until its nix files change" , 640file = sys .stderr , 641 ) 642change = capture (nix ,arguments ,_cache ()/ key ,where ) 643record = { 644"ref" :arguments [0 ], 645"delta" :change , 646"profile" :key , 647"format" :FORMAT , 648 } 649_keep (key ,record ) 650# nix may have written the lock file it locked, a shell file too 651if flake is not None and (written := _key (flake ,attribute ))!= key : 652_keep (written ,record ) 653return change 654return record ["delta" ] 655 656 657# ---- the shim ---- 658 659 660def develop (nix :str ,here :str ,command :Command ,cwd :Path )-> None : 661"""Run a `nix develop ... -c program` in its cached dev shell; return 662when a flag asks for more than the shell's environment.""" 663if command .path != ["develop" ]or not command .program : 664return 665if set (command .flags )- QUIET - {"command" }: 666return 667named = [command .words [index ]for index in command .positional ] 668change = dev_shell (nix ,named [0 ]if named else "." ,cwd ) 669if change is None : 670return 671environment = applied (change ,dict (os .environ )) 672# the shims stay first, ahead of the shell's own programs 673shims = [ 674entry 675for entry in os .environ .get ("PATH" ,"" ).split (os .pathsep ) 676if Path (entry ).parent == Path (here ).parent 677 ] 678rest = environment .get ("PATH" ,"" ).split (os .pathsep ) 679environment ["PATH" ]= os .pathsep .join ( 680 [* shims ,* (entry for entry in rest if entry not in shims )] 681 ) 682found = shutil .which (command .program [0 ],path = environment ["PATH" ]) 683if found is None : 684sys .exit (f" { NOTE } { command . program [ 0 ] } : not found in the dev shell" ) 685os .execve (found ,command .program ,environment ) 686 687 688def main (argv :list [str ])-> NoReturn : 689here ,words = argv [1 ],argv [2 :] 690path = os .environ .get ("PATH" ,"" ).split (os .pathsep ) 691nix = shutil .which ("nix" ,path = os .pathsep .join (e for e in path if e != here )) 692if nix is None : 693sys .exit (f" { NOTE } nix is not on PATH" ) 694cwd = Path .cwd () 695try : 696table = flag_table (nix ) 697except NixError as error : 698flake = locate (cwd ) 699if flake is not None and flake .copied : 700sys .exit (f" { NOTE } { error } \n { NOTE } { REFUSED } " ) 701os .execv (nix , [nix ,* words ]) 702try : 703command = read (words ,table ) 704except Unreadable : 705os .execv (nix , [nix ,* words ])# a flag nix does not know: nix says so 706try : 707develop (nix ,here ,command ,cwd ) 708result = rewritten (command ,cwd ) 709except NixError as error : 710sys .exit (f" { NOTE } { error } " ) 711if result is None : 712os .execv (nix , [nix ,* words ]) 713words ,extra = result 714f" { NOTE } reading this jj workspace from its commit, not copying it: " 716f"nix { ' ' . join ( words ) } " , 717file = sys .stderr , 718 ) 719os .execve (nix , [nix ,* words ], {** os .environ ,** extra }) 720 721 722if __name__ == "__main__" : 723main (sys .argv )