maintainers/scripts/update.nix: refactoring

Get rid of some globals, split main into smaller functions, rename some variables, add typehints.

maintainers/scripts/update.py

@@ -1,3 +1,4 @@
+from typing import Dict, Generator, Tuple, Union
 import argparse
 import contextlib
 import concurrent.futures
@@ -8,28 +9,30 @@ import sys
 import tempfile
 import threading
 
-updates = {}
+updates: Dict[concurrent.futures.Future, Dict] = {}
 
-thread_name_prefix='UpdateScriptThread'
+TempDirs = Dict[str, Tuple[str, str, threading.Lock]]
+
+thread_name_prefix = 'UpdateScriptThread'
 
 def eprint(*args, **kwargs):
     print(*args, file=sys.stderr, **kwargs)
 
-def run_update_script(package, commit):
+def run_update_script(package: Dict, commit: bool, temp_dirs: TempDirs) -> subprocess.CompletedProcess:
+    worktree: Union[None, str] = None
+
     if commit and 'commit' in package['supportedFeatures']:
         thread_name = threading.current_thread().name
         worktree, _branch, lock = temp_dirs[thread_name]
         lock.acquire()
         package['thread'] = thread_name
-    else:
-        worktree = None
 
     eprint(f" - {package['name']}: UPDATING ...")
 
     return subprocess.run(package['updateScript'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, cwd=worktree)
 
 @contextlib.contextmanager
-def make_worktree():
+def make_worktree() -> Generator[Tuple[str, str], None, None]:
     with tempfile.TemporaryDirectory() as wt:
         branch_name = f'update-{os.path.basename(wt)}'
         target_directory = f'{wt}/nixpkgs'
@@ -39,7 +42,40 @@ def make_worktree():
         subprocess.run(['git', 'worktree', 'remove', target_directory], check=True)
         subprocess.run(['git', 'branch', '-D', branch_name], check=True)
 
-def main(max_workers, keep_going, commit, packages):
+def commit_changes(worktree: str, branch: str, execution: subprocess.CompletedProcess) -> None:
+    changes = json.loads(execution.stdout)
+    for change in changes:
+        subprocess.run(['git', 'add'] + change['files'], check=True, cwd=worktree)
+        commit_message = '{attrPath}: {oldVersion} → {newVersion}'.format(**change)
+        subprocess.run(['git', 'commit', '-m', commit_message], check=True, cwd=worktree)
+        subprocess.run(['git', 'cherry-pick', branch], check=True)
+
+def merge_changes(package: Dict, future: concurrent.futures.Future, commit: bool, keep_going: bool, temp_dirs: TempDirs) -> None:
+    try:
+        execution = future.result()
+        if commit and 'commit' in package['supportedFeatures']:
+            thread_name = package['thread']
+            worktree, branch, lock = temp_dirs[thread_name]
+            commit_changes(worktree, branch, execution)
+        eprint(f" - {package['name']}: DONE.")
+    except subprocess.CalledProcessError as e:
+        eprint(f" - {package['name']}: ERROR")
+        eprint()
+        eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+        eprint()
+        eprint(e.stdout.decode('utf-8'))
+        with open(f"{package['pname']}.log", 'wb') as logfile:
+            logfile.write(e.stdout)
+        eprint()
+        eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
+
+        if not keep_going:
+            sys.exit(1)
+    finally:
+        if commit and 'commit' in package['supportedFeatures']:
+            lock.release()
+
+def main(max_workers: int, keep_going: bool, commit: bool, packages: Dict) -> None:
     with open(sys.argv[1]) as f:
         packages = json.load(f)
 
@@ -55,45 +91,18 @@ def main(max_workers, keep_going, commit, packages):
         eprint('Running update for:')
 
         with contextlib.ExitStack() as stack, concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix) as executor:
-            global temp_dirs
-
             if commit:
                 temp_dirs = {f'{thread_name_prefix}_{str(i)}': (*stack.enter_context(make_worktree()), threading.Lock()) for i in range(max_workers)}
+            else:
+                temp_dirs = {}
 
             for package in packages:
-                updates[executor.submit(run_update_script, package, commit)] = package
+                updates[executor.submit(run_update_script, package, commit, temp_dirs)] = package
 
             for future in concurrent.futures.as_completed(updates):
                 package = updates[future]
 
-                try:
+                merge_changes(package, future, commit, keep_going, temp_dirs)
-                    p = future.result()
-                    if commit and 'commit' in package['supportedFeatures']:
-                        thread_name = package['thread']
-                        worktree, branch, lock = temp_dirs[thread_name]
-                        changes = json.loads(p.stdout)
-                        for change in changes:
-                            subprocess.run(['git', 'add'] + change['files'], check=True, cwd=worktree)
-                            commit_message = '{attrPath}: {oldVersion} → {newVersion}'.format(**change)
-                            subprocess.run(['git', 'commit', '-m', commit_message], check=True, cwd=worktree)
-                            subprocess.run(['git', 'cherry-pick', branch], check=True)
-                    eprint(f" - {package['name']}: DONE.")
-                except subprocess.CalledProcessError as e:
-                    eprint(f" - {package['name']}: ERROR")
-                    eprint()
-                    eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
-                    eprint()
-                    eprint(e.stdout.decode('utf-8'))
-                    with open(f"{package['pname']}.log", 'wb') as f:
-                        f.write(e.stdout)
-                    eprint()
-                    eprint(f"--- SHOWING ERROR LOG FOR {package['name']} ----------------------")
-
-                    if not keep_going:
-                        sys.exit(1)
-                finally:
-                    if commit and 'commit' in package['supportedFeatures']:
-                        lock.release()
 
         eprint()
         eprint('Packages updated!')