Files @ 686166c7d5bc
Branch filter:

Location: Diana/src/diana/diana.py

Laman
added more whitespace
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 collectMoves(root):
	node = root
	while len(node.children) > 0:
		b = node.getProp("B")
		w = node.getProp("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, fileName):
		self.fileName = fileName
		self._shortName = "".join(re.split(r'[/\\]', fileName)[-1].split('.')[:-1])
		self._game = go.Go()

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

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

		i = 1
		for k in range(0,len(self._moves), cfg.movesPerDiagram):
			filename = os.path.join(cfg.outputDir, "{0}-{1}".format(self._shortName, i))
			self.createDiagram(k, k+cfg.movesPerDiagram).save(filename, "templ-pleb.svg")
			i += 1

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

	def createDiagram(self, start, end):
		# initialize the diagram
		template = Svg()

		self._setMove(start)

		# draw current state
		for lineNumber, line in enumerate(self._game.board):
			for itemNumber, item in enumerate(line):
				if item != EMPTY:
					template.addStone(itemNumber, lineNumber, "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.keepBroken:
					continue
				else:
					return False

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

		return template

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

	def _setMove(self, k):
		self._game = go.Go()

		blackStones = self._record.root.getProp("AB")
		whiteStones = self._record.root.getProp("AW")
		if blackStones:
			for p in blackStones:
				self._game.board[p.r][p.c] = BLACK
		if whiteStones:
			for p in whiteStones:
				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):
		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
			msg = "illegal move: {0} at {1},{2}".format(self._game.moveCount+1, c, r)
			if cfg.keepBroken:
				print(msg)
			else:
				msg += ". aborted"
				print(msg)
				return False
		return True


def main():
	cfg.parseArgs()
	print("processing:")
	files = cfg.inputFiles[:]

	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))