Refactor generate_gradients.py

tools/generate_gradients.py

@@ -45,19 +45,6 @@ def nums(s):
 	return [int(i) for i in s.split()]
 
 
-p = argparse.ArgumentParser(description=__doc__)
-p.add_argument('gradient', nargs='*', metavar='COLOR', type=color, help='List of colors (either indexes from 8-bit palette or 24-bit RGB in hexadecimal notation)')
-p.add_argument('-n', '--num_items', metavar='INT', type=int, help='Number of items in resulting list', default=101)
-p.add_argument('-N', '--num_output', metavar='INT', type=int, help='Number of characters in sample', default=101)
-p.add_argument('-r', '--range', metavar='V1 V2', type=num2, help='Use this range when outputting scale')
-p.add_argument('-s', '--show', action='store_true', help='If present output gradient sample')
-p.add_argument('-p', '--palette', choices=('16', '256'), help='Use this palette. Defaults to 240-color palette (256 colors without first 16)')
-p.add_argument('-w', '--weights', metavar='INT INT ...', type=nums, help='Adjust weights of colors. Number of weights must be equal to number of colors')
-p.add_argument('-C', '--omit-terminal', action='store_true', help='If present do not compute values for terminal')
-
-args = p.parse_args()
-
-
 def linear_gradient(start_value, stop_value, start_offset, stop_offset, offset):
 	return start_value + ((offset - start_offset) * (stop_value - start_value) / (stop_offset - start_offset))
 
@@ -133,62 +120,91 @@ def dec_scale_generator(num):
 	return r
 
 
-m = args.num_items
+def compute_steps(gradient, weights):
+	maxweight = len(gradient) - 1
+	if weights:
+		weight_sum = sum(weights)
+		norm_weights = [100.0 * weight / weight_sum for weight in weights]
+		steps = [0]
+		for weight in norm_weights:
+			steps.append(steps[-1] + weight)
+		steps.pop(0)
+		steps.pop(0)
+	else:
+		step = m / maxweight
+		steps = [i * step for i in range(1, maxweight + 1)]
+	return steps
 
-maxweight = len(args.gradient) - 1
-if args.weights:
-	weight_sum = sum(args.weights)
-	norm_weights = [100.0 * weight / weight_sum for weight in args.weights]
-	steps = [0]
-	for weight in norm_weights:
-		steps.append(steps[-1] + weight)
-	steps.pop(0)
-	steps.pop(0)
-else:
-	step = m / maxweight
-	steps = [i * step for i in range(1, maxweight + 1)]
 
-data = [(weight, args.gradient[i - 1], args.gradient[i]) for weight, i in zip(steps, range(1, len(args.gradient)))]
-gr_func = generate_gradient_function(data)
-gradient = [gr_func(y) for y in range(0, m)]
 palettes = {
 	'16': (cterm_to_lab[:16], lambda c: c),
 	'256': (cterm_to_lab, lambda c: c),
 	None: (cterm_to_lab[16:], lambda c: c + 16),
 }
-r = [get_rgb(lab) for lab in gradient]
+
-if not args.omit_terminal:
+
-	r2 = [find_color(lab, *palettes[args.palette])[0] for lab in gradient]
+def show_scale(rng, num_output):
-	r3 = [i[0] for i in groupby(r2)]
+	if not rng and num_output >= 32 and (num_output - 1) // 10 >= 4 and (num_output - 1) % 10 == 0:
-print(json.dumps(r))
-if not args.omit_terminal:
-	print(json.dumps(r2))
-	print(json.dumps(r3))
-if args.show:
-	print_colors(args.gradient, args.num_output)
-	print_colors(gradient, args.num_output)
-	if not args.omit_terminal:
-		print_colors(r2, args.num_output)
-		print_colors(r3, args.num_output)
-	if not args.range and args.num_output >= 32 and (args.num_output - 1) // 10 >= 4 and (args.num_output - 1) % 10 == 0:
 		sys.stdout.write('0')
-		sys.stdout.write(''.join(('%*u' % (args.num_output // 10, i) for i in range(10, 101, 10))))
+		sys.stdout.write(''.join(('%*u' % (num_output // 10, i) for i in range(10, 101, 10))))
 		sys.stdout.write('\n')
 	else:
-		if args.range:
+		if rng:
-			vmin, vmax = args.range[1]
+			vmin, vmax = rng[1]
-			isint = args.range[0]
+			isint = rng[0]
 		else:
 			isint = True
 			vmin = 0
 			vmax = 100
 		s = ''
 		lasts = ' ' + str(vmax)
-		while len(s) + len(lasts) < args.num_output:
+		while len(s) + len(lasts) < num_output:
 			curpc = len(s) + 1 if s else 0
-			curval = vmin + curpc * (vmax - vmin) / args.num_output
+			curval = vmin + curpc * (vmax - vmin) / num_output
 			if isint:
 				curval = int(round(curval))
 			s += str(curval) + ' '
 		sys.stdout.write(s[:-1] + lasts + '\n')
-	sys.stdout.write(dec_scale_generator(args.num_output) + '\n')
+	sys.stdout.write(dec_scale_generator(num_output) + '\n')
+
+
+if __name__ == '__main__':
+	p = argparse.ArgumentParser(description=__doc__)
+	p.add_argument('gradient', nargs='*', metavar='COLOR', type=color, help='List of colors (either indexes from 8-bit palette or 24-bit RGB in hexadecimal notation)')
+	p.add_argument('-n', '--num_items', metavar='INT', type=int, help='Number of items in resulting list', default=101)
+	p.add_argument('-N', '--num_output', metavar='INT', type=int, help='Number of characters in sample', default=101)
+	p.add_argument('-r', '--range', metavar='V1 V2', type=num2, help='Use this range when outputting scale')
+	p.add_argument('-s', '--show', action='store_true', help='If present output gradient sample')
+	p.add_argument('-p', '--palette', choices=('16', '256'), help='Use this palette. Defaults to 240-color palette (256 colors without first 16)')
+	p.add_argument('-w', '--weights', metavar='INT INT ...', type=nums, help='Adjust weights of colors. Number of weights must be equal to number of colors')
+	p.add_argument('-C', '--omit-terminal', action='store_true', help='If present do not compute values for terminal')
+
+	args = p.parse_args()
+
+	m = args.num_items
+
+	steps = compute_steps(args.gradient, args.weights)
+
+	data = [(weight, args.gradient[i - 1], args.gradient[i])
+			for weight, i in zip(steps, range(1, len(args.gradient)))]
+	gr_func = generate_gradient_function(data)
+	gradient = [gr_func(y) for y in range(0, m)]
+
+	r = [get_rgb(lab) for lab in gradient]
+	if not args.omit_terminal:
+		r2 = [find_color(lab, *palettes[args.palette])[0] for lab in gradient]
+		r3 = [i[0] for i in groupby(r2)]
+
+	print(json.dumps(r))
+	if not args.omit_terminal:
+		print(json.dumps(r2))
+		print(json.dumps(r3))
+
+	if args.show:
+		print_colors(args.gradient, args.num_output)
+		print_colors(gradient, args.num_output)
+		if not args.omit_terminal:
+			print_colors(r2, args.num_output)
+			print_colors(r3, args.num_output)
+
+		show_scale(args.range, args.num_output)