shrinkwrap-tool 2026.2.1__py3-none-any.whl

shrinkwrap/utils/graph.py

@@ -0,0 +1,263 @@
+# Copyright (c) 2022, Arm Limited.
+# SPDX-License-Identifier: MIT
+
+import graphlib
+import os
+import shutil
+import tempfile
+import shrinkwrap.utils.logger as logger
+import shrinkwrap.utils.label as label
+import shrinkwrap.utils.process as process
+import shrinkwrap.utils.workspace as workspace
+
+
+def _mk_labels(graph):
+	"""
+	Returns a tuple of labels and mask. Each is a two level dictionary,
+	where there is an entry for each config/component that we are building.
+	Mask is initially set to true.
+	"""
+	labels = {}
+	mask = {}
+
+	ts = graphlib.TopologicalSorter(graph)
+	ts.prepare()
+	while ts.is_active():
+		for frag in ts.get_ready():
+			ts.done(frag)
+
+			cfg = frag.config
+			cmp = frag.component
+
+			if cfg is None or cmp is None:
+				continue
+
+			if cfg not in labels:
+				labels[cfg] = {}
+				mask[cfg] = {}
+			if cmp not in labels[cfg]:
+				labels[cfg][cmp] = label.Label()
+				mask[cfg][cmp] = True
+
+	return labels, mask
+
+
+def _mk_label_controller(labels, overdraw):
+	"""
+	Makes a label controller for all the labels in the labels two-level
+	dictionary.
+	"""
+	label_list = []
+	for sl in labels.values():
+		for l in sl.values():
+			label_list.append(l)
+	return label.LabelController(label_list, overdraw=overdraw)
+
+
+def _mk_tag(config, component):
+	"""
+	Makes a fixed-width tag string for a config and component.
+	"""
+	def _clamp(text, max):
+		if len(text) > max:
+			text = text[:max-3] + '...'
+		return text
+
+	config = '' if config is None else config
+	component = '' if component is None else component
+	config = _clamp(config, 10)
+	component = _clamp(component, 14)
+	return '[ {:>10} : {:14} ]'.format(config, component)
+
+
+def _update_labels(labels, mask, config, component, summary):
+	"""
+	Updates all the labels whose mask is True and which match config and
+	component (if specified, then only update for that config/component. If
+	None, then update all configs/components). summary provides the text to
+	update the labels with.
+	"""
+	def iter(labels, mask, key):
+		if key is None:
+			for key in labels:
+				yield key, labels[key], mask[key]
+		else:
+			yield key, labels[key], mask[key]
+
+	for cfg, l0, m0 in iter(labels, mask, config):
+		for cmp, l1, m1 in iter(l0, m0, component):
+			if m1:
+				l1.update(_mk_tag(cfg, cmp) + ' ' + summary)
+
+
+def _run_script(pm, data, script):
+	# Write the script out to a file in a temp directory, and wrap the
+	# directory name and command to run in a Process. Add the Process to the
+	# ProcessManager. On completion, the caller must destroy the directory.
+
+	tmpdir = tempfile.mkdtemp(dir=workspace.build)
+	tmpfilename = os.path.join(tmpdir, 'script.sh')
+	with open(tmpfilename, 'w') as tmpfile:
+		tmpfile.write(script.commands())
+
+	# Start the process asynchronously.
+	pm.add(process.Process(f'bash {tmpfilename}',
+			       False,
+			       (*data, script, tmpdir),
+			       True))
+
+
+def execute(graph, tasks, verbose=False, colorize=True):
+	labels, mask = _mk_labels(graph)
+	lc = _mk_label_controller(labels, not verbose)
+
+	queue = []
+	active = 0
+	log = logger.Logger(27)
+	ts = graphlib.TopologicalSorter(graph)
+	lognum = {}
+	exec_error = None
+
+	def _pump(pm):
+		nonlocal queue
+		nonlocal active
+		nonlocal log
+		nonlocal lognum
+		while len(queue) > 0 and active < tasks:
+			frag = queue.pop()
+			logname = None
+			if frag.config and frag.component:
+				if frag.component not in lognum:
+					lognum[frag.component] = 0
+				logname = os.path.join(workspace.build,
+						'log',
+						frag.config,
+						f'{frag.component}{lognum[frag.component]}.log')
+				os.makedirs(os.path.dirname(logname), exist_ok=True)
+				lognum[frag.component] += 1
+			_update_labels(labels,
+				       mask,
+				       frag.config,
+				       frag.component,
+				       frag.summary + '...')
+			data = (log.alloc_data(str(frag), colorize, False, logname), [])
+			_run_script(pm, data, frag)
+			active += 1
+
+	def _should_log(proc, data, streamid):
+		if verbose or \
+		   (streamid == process.STDERR and \
+		   (not proc.data[2].stderrfilt or \
+		   'warning' in data or 'error' in data)):
+			return True
+		return False
+
+	def _log(pm, proc, data, streamid):
+		logstd = _should_log(proc, data, streamid)
+		if not verbose:
+			proc.data[1].append(data)
+		if logstd:
+			lc.erase()
+		log.log(pm, proc, data, streamid, logstd)
+		if logstd:
+			lc.update()
+
+	def _complete(pm, proc, retcode):
+		nonlocal queue
+		nonlocal active
+		nonlocal ts
+		nonlocal exec_error
+
+		data = proc.data[0]
+		err = proc.data[1]
+		frag = proc.data[2]
+		tmpdir = proc.data[3]
+
+		log.free_data(data)
+		shutil.rmtree(tmpdir)
+
+		if retcode is None:
+			# Forcibly terminated due to errors elsewhere. No need
+			# to do anything further.
+			return
+
+
+		if retcode:
+			# Fatal error. Do not execute any new fragment, only wait for those
+			# that are currently running.
+			exec_error = {
+				"exception": Exception(f"Failed to execute '{frag}'"),
+				"error": ''.join(err),
+			}
+			queue = []
+			state = 'Error'
+		else:
+			state = 'Done' if frag.final else 'Waiting...'
+
+		_update_labels(labels,
+			       mask,
+			       frag.config,
+			       frag.component,
+			       state)
+		if frag.final:
+			mask[frag.config][frag.component] = False
+
+		active -= 1
+		if not exec_error:
+			ts.done(frag)
+			queue.extend(ts.get_ready())
+		_pump(pm)
+
+		lc.update()
+
+	# Initially set all labels to waiting. They will be updated as the
+	# fragments execute.
+	_update_labels(labels, mask, None, None, 'Waiting...')
+
+	# The process manager will run all added processes in the background and
+	# give callbacks whenever there is output available and when each
+	# process terminates. _pump() adds processes to the set.
+	pm = process.ProcessManager(_log, _complete)
+
+	# Fill the queue with all the initial script fragments which do not have
+	# start dependencies.
+	ts.prepare()
+	queue.extend(ts.get_ready())
+
+	# Call _pump() initially to start as many tasks as are allowed.
+	# Then enter the pm.
+	_pump(pm)
+	lc.update()
+	pm.run()
+
+	if exec_error:
+		if not verbose:
+			print('\n== error start ' + ('=' * 65))
+			print(exec_error['error'])
+			print('== error end ' + ('=' * 67) + '\n')
+		raise exec_error['exception']
+
+	# Mark all components as done. This should be a nop since the script
+	# should have indicated if it was the last step for a given
+	# config/component and we would have already set it to done. But this
+	# catches anything that might have slipped through.
+	_update_labels(labels, mask, None, None, 'Done')
+	lc.update()
+
+
+def make_script(graph):
+	# Start the script with the preamble from the first script fragment in
+	# the graph. The preamble for each script is identical and we only need
+	# it once since we are concatenating the fragments together.
+	script = '' + list(graph.keys())[0].preamble() + '\n'
+
+	# Walk the graph, adding each script fragment to the final script
+	# (without its preamble).
+	ts = graphlib.TopologicalSorter(graph)
+	ts.prepare()
+	while ts.is_active():
+		for frag in ts.get_ready():
+			script += frag.commands(False) + '\n'
+			ts.done(frag)
+
+	return script

shrinkwrap/utils/label.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2022, Arm Limited.
+# SPDX-License-Identifier: MIT
+
+import os
+import re
+import sys
+import shrinkwrap.utils.tty as tty
+
+
+class Label:
+	"""
+	Container for a string. Intended to be used with LabelController to
+	display labels (strings of text) on a terminal at a fixed location and
+	be able to update them in place. Instantiate with desired text, and
+	update the text by calling update(). Nothing actually changes on screen
+	until the LabelController does an update().
+	"""
+	def __init__(self, text=''):
+		self.text = text
+		self._prev_text = ''
+		self._lc = None
+
+	def update(self, text=''):
+		self.text = text
+		if self._lc:
+			self._lc._update_pending = True
+
+
+class LabelController:
+	"""
+	Coordinates printing and updating a set of labels at a fixed location on
+	a terminal. The controller manages a list of labels, which are provided
+	at construction and are displayed one under the other on the terminal.
+	update() can be periodically called to redraw the labels if their
+	contents has changed. Only works as long as no other text is printed to
+	the terminal while the labels are active. If the provided file is not
+	backed by a terminal, overdrawing is not performed.
+	"""
+	def __init__(self, labels=[], overdraw=True):
+		self._labels = labels
+		self._overdraw = overdraw
+		self._drawn = False
+		self._cursor_col = 1
+		self._separator = ''
+		self._update_pending = True
+
+		self._out = sys.stdout
+		self._in = sys.stdin
+
+		if overdraw and not self._term_cols():
+			self._overdraw = False
+
+		for label in self._labels:
+			label._lc = self
+
+	def _term_cols(self):
+		try:
+			term_sz = os.get_terminal_size(self._out.fileno())
+			return term_sz.columns
+		except OSError:
+			return None
+
+	def cursor_pos(self):
+		orig = tty.configure(self._in)
+		if orig:
+			try:
+				pos = ''
+				self._out.write('\x1b[6n')
+				self._out.flush()
+				while not (pos := pos + self._in.read(1)).endswith('R'):
+					pass
+				res = re.match(r'.*\[(?P<y>\d*);(?P<x>\d*)R', pos)
+			finally:
+				tty.restore(self._in, orig)
+			if(res):
+				return (int(res.group("x")), int(res.group("y")))
+		assert(False)
+
+	def _move_up(self, line_count, column=1):
+		assert(self._overdraw)
+		self._out.write(f'\033[{line_count}F')
+		if column > 1:
+			self._out.write(f'\033[{column}G')
+
+	def _line_count(self, text, term_cols):
+		assert(self._overdraw)
+		return (len(text) + term_cols - 1) // term_cols
+
+	def update(self):
+		if not self._update_pending:
+			return
+		if self._overdraw:
+			term_cols = self._term_cols()
+			if not self._drawn:
+				self._cursor_col = self.cursor_pos()[0]
+				if self._cursor_col > 1:
+					self._out.write('\n')
+				self._separator = '-' * term_cols
+				self._out.write(self._separator)
+				self._out.write('\n')
+				for l in self._labels:
+					self._out.write(l.text)
+					self._out.write('\n')
+					l._prev_text = l.text
+			else:
+				erase_lines = 0
+				for l in self._labels:
+					erase_lines += self._line_count(l._prev_text, term_cols)
+				self._move_up(erase_lines)
+				for l in self._labels:
+					cc = len(l.text)
+					pcc = len(l._prev_text)
+					self._out.write(l.text)
+					self._out.write(' ' * (pcc - cc))
+					self._out.write('\n')
+					l._prev_text = l.text
+		else:
+			for l in self._labels:
+				if l.text != l._prev_text:
+					self._out.write(l.text)
+					self._out.write('\n')
+					l._prev_text = l.text
+		self._out.flush()
+		self._drawn = True
+		self._update_pending = False
+
+	def erase(self):
+		if not self._overdraw or not self._drawn:
+			return
+
+		term_cols = self._term_cols()
+
+		erase_lines = self._line_count(self._separator, term_cols)
+		for l in self._labels:
+			erase_lines += self._line_count(l._prev_text, term_cols)
+
+		self._move_up(erase_lines)
+
+		self._out.write(' ' * len(self._separator))
+		self._out.write('\n')
+
+		for l in self._labels:
+			pcc = len(l._prev_text)
+			self._out.write(' ' * pcc)
+			self._out.write('\n')
+
+		if self._cursor_col > 1:
+			erase_lines += 1
+		self._move_up(erase_lines, self._cursor_col)
+
+		self._out.flush()
+		self._drawn = False
+		self._update_pending = True

shrinkwrap/utils/logger.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2022, Arm Limited.
+# SPDX-License-Identifier: MIT
+
+import sys
+from collections import namedtuple
+import re
+termcolor = None
+
+
+def _import_termcolor():
+	global termcolor
+	import termcolor as tc
+	termcolor = tc
+
+
+_ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+_colors = ['blue', 'cyan', 'green', 'yellow', 'magenta']
+Data = namedtuple("Data", "id tag color noesc escbuf logfile")
+
+
+class MatchBuf:
+	def __init__(self, match):
+		self._match = match
+		self._buf = ''
+
+	def match(self, data):
+		self._buf += data
+		found = self._buf.find(self._match) >= 0
+		self._buf = self._buf[1 - len(self._match):]
+		return found
+
+
+def splitlines(string):
+	"""
+	Like str.splitlines(True) but preserves '\r'.
+	"""
+	lines = string.split('\n')
+	keepends = [l + '\n' for l in lines[:-1]]
+	if lines[-1] != '':
+		keepends.append(lines[-1])
+	return keepends
+
+
+class Logger:
+	def __init__(self, tag_size):
+		self._tag_size = tag_size
+		self._id_next = 0
+		self._color_next = 0
+		self._prev_id = None
+		self._prev_char = '\n'
+
+	def alloc_data(self, tag, colorize, no_escapes=False, logname=None):
+		"""
+		Returns the object that should be stashed in proc.data[0] when
+		log() is called. Includes the tag for the process and an
+		allocated colour.
+		"""
+		if colorize:
+			_import_termcolor()
+			color = self._color_next
+			self._color_next += 1
+			color = _colors[color % len(_colors)]
+		else:
+			color = None
+
+		id = self._id_next
+		self._id_next += 1
+
+		if isinstance(no_escapes, str):
+			noesc = True
+			escbuf = MatchBuf(no_escapes)
+		else:
+			noesc = no_escapes
+			escbuf = None
+
+		if logname:
+			logfile = open(logname, 'w', buffering=1)
+		else:
+			logfile = None
+
+		return Data(id, tag, color, [noesc], escbuf, logfile)
+
+	def free_data(self, data):
+		if data.logfile:
+			data.logfile.close()
+
+	def log(self, pm, proc, data, streamid, logstd=True):
+		"""
+		Logs text data from one of the processes (FVP or one of its uart
+		terminals) to the terminal. Text is colored and a tag is added
+		on the left to identify the originating process.
+		"""
+		id = proc.data[0].id
+		tag = proc.data[0].tag
+		color = proc.data[0].color
+		noesc = proc.data[0].noesc
+		escbuf = proc.data[0].escbuf
+		logfile = proc.data[0].logfile
+
+		# Write out to file if requested.
+		if logfile:
+			logfile.write(data)
+
+		# Write to stdout if requested.
+		if not logstd:
+			return
+
+		# Remove any ansi escape sequences if requested.
+		if noesc[0]:
+			if escbuf and escbuf.match(data):
+				noesc[0] = False
+			else:
+				data = _ansi_escape.sub('', data)
+				if len(data) == 0:
+					return
+
+		# Make the tag.
+		if len(tag) > self._tag_size:
+			tag = tag[:self._tag_size-3] + '...'
+		if len(tag):
+			tag = f'{{:>{self._tag_size}}}'.format(tag)
+
+		lines = splitlines(data)
+		start = 0
+
+		# If the cursor is not at the start of a new line, then if the
+		# new log is for the same proc that owns the first part of the
+		# line, just continue from there without adding a new tag. If
+		# the first part of the line has a different owner, insert a
+		# newline and add a tag for the new owner.
+		if self._prev_char != '\n':
+			if self._prev_id == id:
+				self.print(lines[0], tag, True, color, end='')
+				start = 1
+			else:
+				self.print('\n', tag, True, color, end='')
+
+		for line in lines[start:]:
+			self.print(line, tag, False, color, end='')
+
+		self._prev_id = id
+		self._prev_char = lines[-1][-1]
+
+		sys.stdout.flush()
+
+	def print(self, text, tag, cont, color=None, on_color=None, attrs=None, **kwargs):
+		# Ensure that any '\r's only rewind to the end of the tag.
+		if len(tag):
+			tag = f'[ {tag} ] '
+		text = text.replace('\r', f'\r{tag}')
+
+		if not cont:
+			self._print(tag, color, on_color, attrs, end='')
+
+		self._print(text, color, on_color, attrs, **kwargs)
+
+	def _print(self, text, color=None, on_color=None, attrs=None, **kwargs):
+		if color:
+			text = termcolor.colored(text, color, on_color, attrs)
+		print(text, **kwargs)