Script to adjust costs in KSP tree.cfg file

tree_change_costs.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function
from tempfile import NamedTemporaryFile
from os.path import basename, dirname, realpath
import os, sys, re, shutil


def parse_tree_costs(src):
	tech_costs = dict()
	for line in iter(src.readline, ''):
		if line.strip() == 'NODE':
			tech_id = tech_cost = None
			line_buff = list()
			for line in iter(src.readline, ''):
				match = re.search(r'^\s*techID\s*=\s*(?P<val>\S+)\s*$', line)
				if match: tech_id = match.group('val')
				match = re.search(r'^\s*cost\s*=\s*(?P<val>\d+)\s*$', line)
				if match: tech_cost = int(match.group('val'))
				line_buff.append(line.rstrip())
				if line.strip() == 'PARTS': break
			if tech_id is None or tech_cost is None:
				raise ValueError( 'Got node without'
					' techID/cost, contents: {!r}'.format('\n'.join(line_buff)) )
			tech_costs[tech_id] = tech_cost
	return tech_costs


def process_tree_costs( src, dst,
		ref=None, cost_min=10, cost_factor=0.5 ):
	tech_costs = dict()
	for line in iter(src.readline, ''):
		dst.write(line)
		if line.strip() == 'NODE':
			tech_id, line_buff = None, list()
			for line in iter(src.readline, ''):
				match = re.search(r'^\s*techID\s*=\s*(?P<val>\S+)\s*$', line)
				if match:
					tech_id = match.group('val')
				match = re.search(r'^\s*cost\s*=\s*(?P<val>\d+)\s*$', line)
				if match:
					s, e = match.span('val')
					val = int(match.group('val'))
					if ref:
						if tech_id is None:
							raise ValueError( 'Got node without techID'
								' before cost, contents: {!r}'.format('\n'.join(line_buff)) )
						if val != ref[tech_id]:
							log.debug( 'Reverting value for'
								' techID "%s": %s -> %s', tech_id, val, ref[tech_id] )
							val = ref[tech_id]
					if val > cost_min: val = max(round(val * cost_factor, 0), 1)
					line = line[:s] + str(int(val)) + line[e:]
				line_buff.append(line.rstrip())
				dst.write(line)
				if line.strip() == 'PARTS': break


def main(args=None):
	import argparse
	parser = argparse.ArgumentParser(
		description='Change tech tree node cost values.')
	parser.add_argument('file_src',
		help='File to edit. In-place, if no file_dst specified.')
	parser.add_argument('file_dst', nargs='?',
		help='File to output results to (optional). "-" for output to stdout.')
	parser.add_argument('-r', '--ref-tree', metavar='path',
		help='Tree to revert any changed costs to before processing them.')
	parser.add_argument('-f', '--cost-factor', metavar='float', type=float, default=0.5,
		help='Factor to apply to all tech node costs in the tree (default: %(default)s).')
	parser.add_argument('-m', '--cost-min', metavar='int', type=int, default=10,
		help='Do not apply --cost-factor to tech that costs'
			' less than specified threshold (default: %(default)s).')
	parser.add_argument('-d', '--debug', action='store_true', help='Verbose operation mode.')
	opts = parser.parse_args(sys.argv[1:] if args is None else args)

	global log
	import logging
	logging.basicConfig(level=logging.DEBUG if opts.debug else logging.WARNING)
	log = logging.getLogger()

	stdout = False
	if not opts.file_dst or opts.file_dst == '-':
		if opts.file_dst == '-': stdout = True
		opts.file_dst = opts.file_src
	src = realpath(opts.file_src)

	ref = None
	if opts.ref_tree:
		with open(opts.ref_tree) as ref:
			ref = parse_tree_costs(ref)

	with open(src) as src,\
			NamedTemporaryFile(
				dir=realpath(dirname(opts.file_dst)),
				prefix=basename(opts.file_dst)+'.',
				delete=False ) as tmp:
		try:
			process_tree_costs( src, tmp, ref=ref,
				cost_min=opts.cost_min, cost_factor=opts.cost_factor )
			if not stdout:
				os.rename(tmp.name, src.name)
			else:
				tmp.seek(0)
				shutil.copyfileobj(tmp, sys.stdout)
		finally:
			try: os.unlink(tmp.name)
			except OSError: pass


if __name__ == '__main__': sys.exit(main())

Usage (linux zsh console)

% ./tree_change_costs.py -h
usage: tree_change_costs.py [-h] [-r path] [-f float] [-m int] [-d]
                            file_src [file_dst]

Change tech tree node cost values.

positional arguments:
  file_src              File to edit. In-place, if no file_dst specified.
  file_dst              File to output results to (optional). "-" for output
                        to stdout.

optional arguments:
  -h, --help            show this help message and exit
  -r path, --ref-tree path
                        Tree to revert any changed costs to before processing
                        them.
  -f float, --cost-factor float
                        Factor to apply to all tech node costs in the tree
                        (default: 0.5).
  -m int, --cost-min int
                        Do not apply --cost-factor to tech that costs less
                        than specified threshold (default: 10).
  -d, --debug           Verbose operation mode.

% ./tree_change_costs.py tree.cfg
%