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 col2num(column, board_size):
  45         a, o, z = map(ord, ['a', column, 'z'])
  46         if a <= o <= z:
  47                 return a + board_size - o
  48         raise Exception( "Wrong column character: '%s'"%(column,) )
  49
  50 def is_pass_move(coord, board_size):
  51         # the pass move is represented either by [] ( = empty coord )
  52         # OR by [tt] (for boards <= 19 only)
  53         return len(coord) == 0 or ( board_size <= 19 and coord == 'tt' )
  54
  55 def process_gametree(gametree, fout):
  56         # cursor for tree traversal
  57         c = gametree.cursor()
  58         # first node is the header
  59         header = c.node
  60
  61         handicap = get_atr(header, 'HA')
  62         board_size = int(get_atr(header, 'SZ'))
  63         komi = get_atr(header, 'KM')
  64         player_next, player_other = "B", "W"
  65
  66         print >>fout, "boardsize", board_size
  67         print >>fout, "clear_board"
  68         if komi:
  69                 print >>fout, "komi", komi
  70         if handicap and handicap != '0':
  71                 print >>fout, "fixed_handicap", handicap
  72                 player_next, player_other = player_other, player_next
  73
  74         def print_game_step(coord):
  75                 if is_pass_move(coord, board_size):
  76                         print >>fout, "play", player_next, "pass"
  77                 else:
  78                         x, y = coord
  79                         # The reason for this incredibly weird thing is that
  80                         # the GTP protocol excludes `i` in the coordinates
  81                         # (do you see any purpose in this??)
  82                         if x >= 'i':
  83                                 x = chr(ord(x)+1)
  84                         y = str(col2num(y, board_size))
  85                         print >>fout, "play", player_next, x+y
  86
  87         movenum = 0
  88         # walk the game tree forward
  89         while 1:
  90                 # sgf2gtp.pl ignores n = 0
  91                 if c.atEnd or (args['n'] and movenum >= args['n']):
  92                         break
  93                 c.next()
  94                 movenum += 1
  95
  96                 coord = get_atr(c.node, player_next)
  97                 if coord != None:
  98                         print_game_step(coord)
  99                 else:
 100                         # MAYBE white started?
 101                         # or one of the players plays two time in a row
 102                         player_next, player_other = player_other, player_next
 103                         coord = get_atr(c.node, player_next)
 104                         if coord != None:
 105                                 print_game_step(coord)
 106                         else:
 107                                 # TODO handle weird sgf files better
 108                                 raise UnknownNode
 109
 110                 player_next, player_other = player_other, player_next
 111
 112         if args['g']:
 113                 print >>fout, "genmove", player_next
 114
 115 def process_sgf_file(fin, fout):
 116         sgfdata = fin.read()
 117         col = SGFParser(sgfdata).parse()
 118
 119         for gametree in col:
 120                 try:
 121                         process_gametree(gametree, fout)
 122                 except UnknownNode:
 123                         # Try next game tree in this file
 124                         if DEBUG:
 125                                 print >>sys.stderr, "Unknown Node"
 126                         continue
 127
 128 if __name__ == "__main__":
 129         if not len(args['FILENAMES']):
 130                 process_sgf_file(sys.stdin, sys.stdout)
 131         else:
 132                 for in_filename in args['FILENAMES']:
 133                         if args['stdout_only']:
 134                                 fout = sys.stdout
 135                         else:
 136                                 if re.search('sgf$', in_filename):
 137                                         filename_base = in_filename[:-3]
 138                                 else:
 139                                         filename_base = in_filename
 140                                 # Save the .gtp file
 141                                 out_filename = filename_base + 'gtp'
 142
 143                                 fout = open(out_filename, 'w')
 144
 145                         fin = open(in_filename, 'r')
 146
 147                         process_sgf_file(fin, fout)
 148
 149                         fin.close()
 150                         if not args['stdout_only']:
 151                                 fout.close()
 152