from .node import Node from . import skip_whitespace, ParserError from .game_record import GameRecord class Collection: """Game collection type, a list of game trees.""" def __init__(self, s): self.game_trees = [] i = skip_whitespace(s, 0) if i >= len(s): return elif not GameTree.fits(s, i): raise ParserError("expected a GameTree starting with '('", s, i) while GameTree.fits(s, i): (i, x) = GameTree.create(s, i) self.game_trees.append(x) if i < len(s): raise ParserError("expected EOF", s, i) def list_games(self): """:rtype: Iterator[GameRecord]""" for tree in self.game_trees: for game in tree.list_games(): yield game class GameTree: def __init__(self): self.nodes = [] self.branches = [] @staticmethod def fits(s, i): """Decide if a GameTree fits at the suggested string location. :param str s: an SGF string :param int i: a string location :return: True if there might be a GameTree at the location, False if not""" return i < len(s) and s[i] == "(" @staticmethod def create(s, start): """Create a GameTree from the string and move the current position pointer. :param str s: an SGF string :param int start: a location in s :return: the resulting tree and a pointer to its end :rtype: (int, GameTree)""" assert GameTree.fits(s, start) res = GameTree() i = skip_whitespace(s, start + 1) if not Node.fits(s, i): raise ParserError("expected a Node starting with ';'", s, i) y = None while Node.fits(s, i): (i, x) = Node.create(s, i) res.nodes.append(x) if y: y.add_child(x) x.parent = y y = x i = skip_whitespace(s, i) while GameTree.fits(s, i): (i, x) = GameTree.create(s, i) res.branches.append(x) subroot = x.get_node(0) subroot.parent = y if y: y.add_child(subroot) i = skip_whitespace(s, i) if i >= len(s) or s[i] != ")": raise ParserError("expected end of a GameTree marked by ')'", s, i) i = skip_whitespace(s, i + 1) return (i, res) def list_games(self): """Expand multiple games into distinct GameTrees and yield each.""" if len(self.nodes) == 0: return None for node in self.nodes[0].list_gi_nodes(): yield GameRecord(self._build_subtree(node)) def get_node(self, i): if 0 <= i < len(self.nodes): return self.nodes[i] return None def _build_subtree(self, seed_node): """Create and return a new game tree containing the provided Node. Ancestor nodes are copied, descendants are moved from the seed_node.""" node = seed_node.copy() node.set_children(seed_node.children) seed_node.children = [] while node.parent: new_node = node.parent.copy() new_node.add_child(node) node = new_node return node