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