17f89667b3
Get rid of some globals, split main into smaller functions, rename some variables, add typehints.
129 lines
5.1 KiB
Python
129 lines
5.1 KiB
Python
from typing import Dict, Generator, Tuple, Union
|
|
import argparse
|
|
import contextlib
|
|
import concurrent.futures
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
|
|
updates: Dict[concurrent.futures.Future, Dict] = {}
|
|
|
|
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: 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
|
|
|
|
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() -> Generator[Tuple[str, str], None, None]:
|
|
with tempfile.TemporaryDirectory() as wt:
|
|
branch_name = f'update-{os.path.basename(wt)}'
|
|
target_directory = f'{wt}/nixpkgs'
|
|
|
|
subprocess.run(['git', 'worktree', 'add', '-b', branch_name, target_directory], check=True)
|
|
yield (target_directory, branch_name)
|
|
subprocess.run(['git', 'worktree', 'remove', target_directory], check=True)
|
|
subprocess.run(['git', 'branch', '-D', branch_name], check=True)
|
|
|
|
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)
|
|
|
|
eprint()
|
|
eprint('Going to be running update for following packages:')
|
|
for package in packages:
|
|
eprint(f" - {package['name']}")
|
|
eprint()
|
|
|
|
confirm = input('Press Enter key to continue...')
|
|
if confirm == '':
|
|
eprint()
|
|
eprint('Running update for:')
|
|
|
|
with contextlib.ExitStack() as stack, concurrent.futures.ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix) as executor:
|
|
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, temp_dirs)] = package
|
|
|
|
for future in concurrent.futures.as_completed(updates):
|
|
package = updates[future]
|
|
|
|
merge_changes(package, future, commit, keep_going, temp_dirs)
|
|
|
|
eprint()
|
|
eprint('Packages updated!')
|
|
sys.exit()
|
|
else:
|
|
eprint('Aborting!')
|
|
sys.exit(130)
|
|
|
|
parser = argparse.ArgumentParser(description='Update packages')
|
|
parser.add_argument('--max-workers', '-j', dest='max_workers', type=int, help='Number of updates to run concurrently', nargs='?', default=4)
|
|
parser.add_argument('--keep-going', '-k', dest='keep_going', action='store_true', help='Do not stop after first failure')
|
|
parser.add_argument('--commit', '-c', dest='commit', action='store_true', help='Commit the changes')
|
|
parser.add_argument('packages', help='JSON file containing the list of package names and their update scripts')
|
|
|
|
if __name__ == '__main__':
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
main(args.max_workers, args.keep_going, args.commit, args.packages)
|
|
except (KeyboardInterrupt, SystemExit) as e:
|
|
for update in updates:
|
|
update.cancel()
|
|
|
|
sys.exit(e.code if isinstance(e, SystemExit) else 130)
|