1#!/usr/bin/env python 2 3import json 4import os 5import re 6 7import pexpect 8from pexpect.popen_spawn import PopenSpawn 9 10from pychess.Players.ProtocolEngine import TIME_OUT_SECOND 11 12PGN_HEADERS_REGEX = re.compile(r"\[([A-Za-z0-9_]+)\s+\"(.*)\"\]") 13 14 15class Scoutfish: 16 def __init__(self, engine=''): 17 if not engine: 18 engine = './scoutfish' 19 self.p = PopenSpawn(engine, timeout=TIME_OUT_SECOND, encoding="utf-8") 20 self.wait_ready() 21 self.pgn = '' 22 self.db = '' 23 24 def wait_ready(self): 25 self.p.sendline('isready') 26 self.p.expect(u'readyok') 27 28 def open(self, pgn): 29 '''Open a PGN file and create an index if not exsisting''' 30 if not os.path.isfile(pgn): 31 raise NameError("File {} does not exsist".format(pgn)) 32 pgn = os.path.normcase(pgn) 33 self.pgn = pgn 34 self.db = os.path.splitext(pgn)[0] + '.scout' 35 if not os.path.isfile(self.db): 36 result = self.make() 37 self.db = result['DB file'] 38 39 def close(self): 40 '''Terminate scoutfish. Not really needed: engine will terminate as 41 soon as pipe is closed, i.e. when we exit.''' 42 self.p.sendline('quit') 43 self.p.expect(pexpect.EOF) 44 self.pgn = '' 45 self.db = '' 46 47 def make(self): 48 '''Make an index out of a pgn file. Normally called by open()''' 49 if not self.pgn: 50 raise NameError("Unknown DB, first open a PGN file") 51 cmd = 'make ' + self.pgn 52 self.p.sendline(cmd) 53 self.wait_ready() 54 s = '{' + self.p.before.split('{')[1] 55 s = s.replace('\\', r'\\') # Escape Windows's path delimiter 56 result = json.loads(s) 57 self.p.before = '' 58 return result 59 60 def setoption(self, name, value): 61 '''Set an option value, like threads number''' 62 cmd = "setoption name {} value {}".format(name, value) 63 self.p.sendline(cmd) 64 self.wait_ready() 65 66 def scout(self, q): 67 '''Run query defined by 'q' dict. Result will be a dict too''' 68 if not self.db: 69 raise NameError("Unknown DB, first open a PGN file") 70 j = json.dumps(q) 71 cmd = "scout {} {}".format(self.db, j) 72 self.p.sendline(cmd) 73 self.wait_ready() 74 result = json.loads(self.p.before) 75 self.p.before = '' 76 return result 77 78 def scout_raw(self, q): 79 '''Run query defined by 'q' dict. Result will be full output''' 80 if not self.db: 81 raise NameError("Unknown DB, first open a PGN file") 82 j = json.dumps(q) 83 cmd = "scout {} {}".format(self.db, j) 84 self.p.sendline(cmd) 85 self.wait_ready() 86 result = self.p.before 87 self.p.before = '' 88 return result 89 90 def get_games(self, matches): 91 '''Retrieve the PGN games specified in the offset list. Games are 92 added to each list item with a 'pgn' key''' 93 if not self.pgn: 94 raise NameError("Unknown DB, first open a PGN file") 95 with open(self.pgn, "rU") as f: 96 for match in matches: 97 f.seek(match['ofs']) 98 game = '' 99 for line in f: 100 if line.startswith('[Event "'): 101 if game: 102 break # Second one, start of next game 103 else: 104 game = line # First occurence 105 elif game: 106 game += line 107 match['pgn'] = game.strip() 108 return matches 109 110 def get_header(self, pgn): 111 '''Return a dict with just header information out of a pgn game. The 112 pgn tags are supposed to be consecutive''' 113 header = {} 114 for line in pgn.splitlines(): 115 line = line.strip() 116 if line.startswith('[') and line.endswith(']'): 117 tag_match = PGN_HEADERS_REGEX.match(line) 118 if tag_match: 119 header[tag_match.group(1)] = tag_match.group(2) 120 else: 121 break 122 return header 123 124 def get_game_headers(self, matches): 125 '''Return a list of headers out of a list of pgn games. It is defined 126 to be compatible with the return value of get_games()''' 127 headers = [] 128 for match in matches: 129 pgn = match['pgn'] 130 h = self.get_header(pgn) 131 headers.append(h) 132 return headers 133