Files @ 966ee650fabf
Branch filter:

Location: Diana/src/diana/diana.py

Laman
updated comments
import os
import re

from . import config as cfg
from . import go
from .go import BLACK, WHITE, EMPTY
from .sgfparser import ParserError
from .sgfparser.collection import Collection
from .drawer.svg import Svg
from .drawer.tikz import Tikz


def collect_moves(root):
	"""Walk the primary variation and collect the move coordinates.

	:param Node root: a node serving as a root for the tree traversal
	:return: sequence of move colors and coordinates (col, row)
	:rtype: Iterator[("b" or "w", (int, int))]"""
	node = root
	while len(node.children) > 0:
		b = node.get_prop("B")
		w = node.get_prop("W")
		if b is not None:
			yield ("b", b)
		elif w is not None:
			yield ("w", w)
		# else: yield None # !! not really robust

		node = node.children[0]


class SourceFile:
	def __init__(self, file_name):
		self.file_name = file_name
		self._short_name = "".join(re.split(r'[/\\]', file_name)[-1].split('.')[:-1])
		self._game = go.Go()

		with open(self.file_name, 'r', encoding=cfg.encoding) as f:
			games = Collection(f.read()).list_games()
		self._record = list(games)[0]
		self._moves = list(collect_moves(self._record.root))

	def process(self):
		print("{0}... ".format(self.file_name), end="")

		i = 1
		for k in range(0, len(self._moves), cfg.moves_per_diagram):
			filename = os.path.join(cfg.output_dir, "{0}-{1}".format(self._short_name, i))
			self.create_diagram(k, k + cfg.moves_per_diagram).save(filename, "templ-pleb.svg")
			i += 1

		info_str = """{GN}
B: {PB} {BR}
W: {PW} {WR}
{DT}
{RE}""".format(**self.fetch_game_info(["GN", "PB", "BR", "PW", "WR", "DT", "RE"], ""))
		notes = open(os.path.join(cfg.output_dir, "{0}.txt".format(self._short_name)), 'w')
		notes.write(info_str)
		notes.close()
		print("done")

	def create_diagram(self, start, end):
		"""Create and return a diagram drawer instance.

		:param int start: the initial move
		:param int end: the first omitted move, similar to sequence slices
		:return: a diagram instance
		:rtype: Drawer"""

		# initialize the diagram
		template = Svg()

		self._set_move(start)

		# draw current state
		for (line_number, line) in enumerate(self._game.board):
			for (item_number, item) in enumerate(line):
				if item != EMPTY:
					template.add_stone(item_number, line_number, "b" if item == BLACK else "w")

		# draw the moves
		for k in range(start, end):
			if k >= len(self._moves):
				break

			color, move = self._moves[k]
			if move == tuple():
				template.overlays.append((k, "pass"))
				continue
			else:
				(c, r) = move

			if not self._move(color, c, r):
				if cfg.keep_broken:
					continue
				else:
					return False

			# draw the move
			template.add_move(c, r, color, k + 1)

		return template

	def fetch_game_info(self, field_names, default=None):
		return {k: self._record.get(k, default) for k in field_names}

	def _set_move(self, k):
		"""Rewind the internal game state to move k.

		:param int k: the move number"""
		self._game = go.Go()

		black_stones = self._record.root.get_prop("AB")
		white_stones = self._record.root.get_prop("AW")
		if black_stones:
			for p in black_stones:
				self._game.board[p.r][p.c] = BLACK
		if white_stones:
			for p in white_stones:
				self._game.board[p.r][p.c] = WHITE

		for i in range(k):
			(color, move) = self._moves[i]
			if move == tuple():
				continue  # pass
			self._move(color, *move)

	def _move(self, color, c, r):
		"""Make a single move.

		:param str color: "b" or "w"
		:param int c: column
		:param int r: row
		:return: True if we can continue with another move, False on abort"""
		if not self._game.move(BLACK if color=='b' else WHITE, c, r):
			# we do not honor http://red-bean.com/sgf/ff5/m_vs_ax.htm at the moment
			# we accept and process only legal moves
			msg = "illegal move: {0} at {1},{2}".format(self._game.move_count+1, c, r)
			if cfg.keep_broken:
				print(msg)
			else:
				msg += ". aborted"
				print(msg)
				return False
		return True


def main():
	cfg.parse_args()
	print("processing:")
	files = cfg.input_files[:]

	for item in files:
		if os.path.isfile(item):
			try:
				f = SourceFile(item)
				f.process()
			except ParserError as e:
				print("Couldn't parse {0}, following error occured: {1}".format(item, e))
		elif os.path.isdir(item):
			files += [os.path.join(item,child) for child in os.listdir(item)]
			print("contents of the '{0}' directory added to the queue".format(item))
		else:
			print("the '{0}' path could not be resolved to either a file nor a directory".format(item))