tools/sgf2gtp.py

   1 #! /usr/bin/env python
   2
   3 import sys
   4 import argparse
   5 import re
   6
   7 from sgflib import SGFParser
   8
   9 DEBUG = False
  10
  11 parser = argparse.ArgumentParser( formatter_class=argparse.RawDescriptionHelpFormatter,
  12 description="""
  13 This script converts SGF files to GTP format so that you can feed them
  14 to Pachi, insert genmove at the right places etc. Might not work on
  15 obscure SGF files.
  16
  17 When called with FILENAMES argument, it will create according output
  18 files with .gtp extension instead of .sgf, unless --no-gtp option is
  19 specified.
  20
  21 Otherwise the games are read from standard input. Note that the stdin
  22 in this case is read in at once, so for very large collections, it is
  23 better to run this script separately for each sgf file.
  24
  25 example:
  26     cat *.sgf | %s -g -n 5
  27 """%(sys.argv[0]))
  28
  29 parser.add_argument('FILENAMES', help='List of sgf games to process.', nargs='*', default=[])
  30 parser.add_argument('-g', help='Automatically append genmove command for the other color.', action='store_true')
  31 parser.add_argument('-n', help='Output at most first MOVENUM moves.', metavar='MOVENUM', type=int, default=10**10)
  32 parser.add_argument('--stdout-only', help='Do not create the .gtp files from FILENAMES, print everything to stdout.', action='store_true')
  33 args = vars(parser.parse_args())
  34
  35 class UnknownNode(Exception):
  36     pass
  37
  38 def get_atr(node, atr):
  39     try:
  40         return node.data[atr].data[0]
  41     except KeyError:
  42         return None
  43
  44 def get_setup(node, atr):
  45     try:
  46         return node.data[atr].data[:]
  47     except KeyError:
  48         return None
  49
  50 def col2num(column, board_size):
  51     a, o, z = map(ord, ['a', column, 'z'])
  52     if a <= o <= z:
  53         return a + board_size - o
  54     raise Exception( "Wrong column character: '%s'"%(column,) )
  55
  56 def is_pass_move(coord, board_size):
  57     # the pass move is represented either by [] ( = empty coord )
  58     # OR by [tt] (for boards <= 19 only)
  59     return len(coord) == 0 or ( board_size <= 19 and coord == 'tt' )
  60
  61 def process_gametree(gametree, fout):
  62     # cursor for tree traversal
  63     c = gametree.cursor()
  64     # first node is the header
  65     header = c.node
  66
  67     handicap = get_atr(header, 'HA')
  68     board_size = int(get_atr(header, 'SZ') or 19)
  69     komi = get_atr(header, 'KM')
  70     player_next, player_other = "B", "W"
  71     setup_black = get_setup(header, 'AB')
  72     setup_white = get_setup(header, 'AW')
  73
  74     print >>fout, "boardsize", board_size
  75     print >>fout, "clear_board"
  76     if komi:
  77         print >>fout, "komi", komi
  78     if handicap and handicap != '0':
  79         print >>fout, "fixed_handicap", handicap
  80         player_next, player_other = player_other, player_next
  81     if setup_black:
  82         for item in setup_black:
  83             x, y = item
  84             if x >= 'i':
  85                 x = chr(ord(x)+1)
  86             y = str(col2num(y, board_size))
  87             print >>fout, "play B", x+y
  88     if setup_white:
  89         for item in setup_white:
  90             x, y = item
  91             if x >= 'i':
  92                 x = chr(ord(x)+1)
  93             y = str(col2num(y, board_size))
  94             print >>fout, "play W", x+y
  95
  96     def print_game_step(coord):
  97         if is_pass_move(coord, board_size):
  98             print >>fout, "play", player_next, "pass"
  99         else:
 100             x, y = coord
 101             # The reason for this incredibly weird thing is that
 102             # the GTP protocol excludes `i` in the coordinates
 103             # (do you see any purpose in this??)
 104             if x >= 'i':
 105                 x = chr(ord(x)+1)
 106             y = str(col2num(y, board_size))
 107             print >>fout, "play", player_next, x+y
 108
 109     movenum = 0
 110     # walk the game tree forward
 111     while 1:
 112         # sgf2gtp.pl ignores n = 0
 113         if c.atEnd or (args['n'] and movenum >= args['n']):
 114             break
 115         c.next()
 116         movenum += 1
 117
 118         coord = get_atr(c.node, player_next)
 119         if coord != None:
 120             print_game_step(coord)
 121         else:
 122             # MAYBE white started?
 123             # or one of the players plays two time in a row
 124             player_next, player_other = player_other, player_next
 125             coord = get_atr(c.node, player_next)
 126             if coord != None:
 127                 print_game_step(coord)
 128             else:
 129                 # TODO handle weird sgf files better
 130                 raise UnknownNode
 131
 132         player_next, player_other = player_other, player_next
 133
 134     if args['g']:
 135         print >>fout, "genmove", player_next
 136
 137 def process_sgf_file(fin, fout):
 138     sgfdata = fin.read()
 139     col = SGFParser(sgfdata).parse()
 140
 141     for gametree in col:
 142         try:
 143             process_gametree(gametree, fout)
 144         except UnknownNode:
 145             # Try next game tree in this file
 146             if DEBUG:
 147                 print >>sys.stderr, "Unknown Node"
 148             continue
 149
 150 if __name__ == "__main__":
 151     if not len(args['FILENAMES']):
 152         process_sgf_file(sys.stdin, sys.stdout)
 153     else:
 154         for in_filename in args['FILENAMES']:
 155             if args['stdout_only']:
 156                 fout = sys.stdout
 157             else:
 158                 if re.search('sgf$', in_filename):
 159                     filename_base = in_filename[:-3]
 160                 else:
 161                     filename_base = in_filename
 162                 # Save the .gtp file
 163                 out_filename = filename_base + 'gtp'
 164
 165                 fout = open(out_filename, 'w')
 166
 167             fin = open(in_filename, 'r')
 168
 169             process_sgf_file(fin, fout)
 170
 171             fin.close()
 172             if not args['stdout_only']:
 173                 fout.close()
 174