1818from _pytask .node_protocols import PTask
1919from _pytask .node_protocols import PTaskWithPath
2020from _pytask .nodes import PythonNode
21+ from _pytask .outcomes import ExitCode
2122from _pytask .pluginmanager import hookimpl
2223
2324if TYPE_CHECKING :
2425 from _pytask .session import Session
2526
26- CURRENT_LOCKFILE_VERSION = "1.0 "
27+ CURRENT_LOCKFILE_VERSION = "1"
2728
2829
2930class LockfileError (Exception ):
@@ -34,20 +35,11 @@ class LockfileVersionError(LockfileError):
3435 """Raised when a lockfile version is not supported."""
3536
3637
37- class _State (msgspec .Struct ):
38- value : str
39-
40-
41- class _NodeEntry (msgspec .Struct ):
42- id : str
43- state : _State
44-
45-
4638class _TaskEntry (msgspec .Struct ):
4739 id : str
48- state : _State
49- depends_on : list [ _NodeEntry ] = msgspec .field (default_factory = list )
50- produces : list [ _NodeEntry ] = msgspec .field (default_factory = list )
40+ state : str
41+ depends_on : dict [ str , str ] = msgspec .field (default_factory = dict )
42+ produces : dict [ str , str ] = msgspec .field (default_factory = dict )
5143
5244
5345class _Lockfile (msgspec .Struct , forbid_unknown_fields = False ):
@@ -109,28 +101,25 @@ def read_lockfile(path: Path) -> _Lockfile | None:
109101 msg = "Lockfile is missing 'lock-version'."
110102 raise LockfileError (msg )
111103
112- if Version (version ) > Version (CURRENT_LOCKFILE_VERSION ):
104+ if Version (version ) != Version (CURRENT_LOCKFILE_VERSION ):
113105 msg = (
114106 f"Unsupported lock-version { version !r} . "
115107 f"Current version is { CURRENT_LOCKFILE_VERSION } ."
116108 )
117109 raise LockfileVersionError (msg )
118110
119- lockfile = msgspec .toml .decode (path .read_bytes (), type = _Lockfile )
120-
121- if Version (version ) < Version (CURRENT_LOCKFILE_VERSION ):
122- lockfile = _Lockfile (
123- lock_version = CURRENT_LOCKFILE_VERSION ,
124- task = lockfile .task ,
125- )
126- return lockfile
111+ try :
112+ return msgspec .toml .decode (path .read_bytes (), type = _Lockfile )
113+ except msgspec .DecodeError :
114+ msg = "Lockfile has invalid format."
115+ raise LockfileError (msg ) from None
127116
128117
129118def _normalize_lockfile (lockfile : _Lockfile ) -> _Lockfile :
130119 tasks = []
131120 for task in sorted (lockfile .task , key = lambda entry : entry .id ):
132- depends_on = sorted ( task .depends_on , key = lambda entry : entry . id )
133- produces = sorted ( task .produces , key = lambda entry : entry . id )
121+ depends_on = { key : task .depends_on [ key ] for key in sorted ( task . depends_on )}
122+ produces = { key : task .produces [ key ] for key in sorted ( task . produces )}
134123 tasks .append (
135124 _TaskEntry (
136125 id = task .id ,
@@ -159,7 +148,7 @@ def _build_task_entry(session: Session, task: PTask, root: Path) -> _TaskEntry |
159148 predecessors = set (dag .predecessors (task .signature ))
160149 successors = set (dag .successors (task .signature ))
161150
162- depends_on = []
151+ depends_on : dict [ str , str ] = {}
163152 for node_signature in predecessors :
164153 node = (
165154 dag .nodes [node_signature ].get ("task" ) or dag .nodes [node_signature ]["node" ]
@@ -174,9 +163,9 @@ def _build_task_entry(session: Session, task: PTask, root: Path) -> _TaskEntry |
174163 if isinstance (node , PTask )
175164 else build_portable_node_id (node , root )
176165 )
177- depends_on . append ( _NodeEntry ( id = node_id , state = _State ( state )))
166+ depends_on [ node_id ] = state
178167
179- produces = []
168+ produces : dict [ str , str ] = {}
180169 for node_signature in successors :
181170 node = (
182171 dag .nodes [node_signature ].get ("task" ) or dag .nodes [node_signature ]["node" ]
@@ -191,11 +180,11 @@ def _build_task_entry(session: Session, task: PTask, root: Path) -> _TaskEntry |
191180 if isinstance (node , PTask )
192181 else build_portable_node_id (node , root )
193182 )
194- produces . append ( _NodeEntry ( id = node_id , state = _State ( state )))
183+ produces [ node_id ] = state
195184
196185 return _TaskEntry (
197186 id = build_portable_task_id (task , root ),
198- state = _State ( task_state ) ,
187+ state = task_state ,
199188 depends_on = depends_on ,
200189 produces = produces ,
201190 )
@@ -238,9 +227,7 @@ def _rebuild_indexes(self) -> None:
238227 self ._task_index = {task .id : task for task in self .lockfile .task }
239228 self ._node_index = {}
240229 for task in self .lockfile .task :
241- nodes = {}
242- for entry in task .depends_on + task .produces :
243- nodes [entry .id ] = entry .state .value
230+ nodes = {** task .depends_on , ** task .produces }
244231 self ._node_index [task .id ] = nodes
245232
246233 def get_task_entry (self , task_id : str ) -> _TaskEntry | None :
@@ -261,10 +248,42 @@ def update_task(self, session: Session, task: PTask) -> None:
261248 self ._rebuild_indexes ()
262249 write_lockfile (self .path , self .lockfile )
263250
251+ def rebuild_from_session (self , session : Session ) -> None :
252+ if session .dag is None :
253+ return
254+ tasks = []
255+ for task in session .tasks :
256+ entry = _build_task_entry (session , task , self .root )
257+ if entry is not None :
258+ tasks .append (entry )
259+ self .lockfile = _Lockfile (
260+ lock_version = CURRENT_LOCKFILE_VERSION ,
261+ task = tasks ,
262+ )
263+ self ._rebuild_indexes ()
264+ write_lockfile (self .path , self .lockfile )
265+
264266
265267@hookimpl
266268def pytask_post_parse (config : dict [str , Any ]) -> None :
267269 """Initialize the lockfile state."""
268270 path = config ["root" ] / "pytask.lock"
269271 config ["lockfile_path" ] = path
270272 config ["lockfile_state" ] = LockfileState .from_path (path , config ["root" ])
273+
274+
275+ @hookimpl
276+ def pytask_unconfigure (session : Session ) -> None :
277+ """Optionally rewrite the lockfile to drop stale entries."""
278+ if session .config .get ("command" ) != "build" :
279+ return
280+ if not session .config .get ("clean_lockfile" ):
281+ return
282+ if session .config .get ("dry_run" ):
283+ return
284+ if session .exit_code != ExitCode .OK :
285+ return
286+ lockfile_state = session .config .get ("lockfile_state" )
287+ if lockfile_state is None :
288+ return
289+ lockfile_state .rebuild_from_session (session )
0 commit comments