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, output_format): self.file_name = file_name self._short_name = "".join(re.split(r'[/\\]', file_name)[-1].split('.')[:-1]) self._game = go.Go() self.drawer = Tikz if output_format == "tikz" else Svg 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) 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 = self.drawer(start) 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, cfg.output_format) 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))