# -*- coding: utf-8 -*-
import collections
import os
import re
import subprocess
import zipfile
from gi.repository import GLib
from sqlalchemy import select, func
from sqlalchemy.exc import SQLAlchemyError
from pychess.Utils.const import NORMALCHESS, RUNNING, DRAW, WHITEWON, BLACKWON
from pychess.Variants import name2variant
from pychess.System.Log import log
from pychess.System import download_file
from pychess.System.protoopen import protoopen, protosave, PGN_ENCODING
from pychess.Database.model import event, site, player, game, annotator, tag_game, source
# from pychess.System import profile_me
# Editable (on game info dialog) tags
dedicated_tags = ('Event', 'Site', 'Date', 'Round', 'White', 'Black', 'WhiteElo', 'BlackElo')
# Other tags stored in game table
other_game_tags = ('Result', 'SetUp', 'FEN', 'ECO', 'Variant', 'PlyCount', 'Annotator', 'offset', 'offset8')
TAG_REGEX = re.compile(r"\[([a-zA-Z0-9_]+)\s+\"(.*)\"\]")
GAME, EVENT, SITE, PLAYER, ANNOTATOR, SOURCE, STAT = range(7)
removeDic = {
ord(u"'"): None,
ord(u","): None,
ord(u"."): None,
ord(u"-"): None,
ord(u" "): None,
}
pgn2Const = {"*": RUNNING,
"?": RUNNING,
"1/2-1/2": DRAW,
"1/2": DRAW,
"1-0": WHITEWON,
"0-1": BLACKWON}
class PgnImport():
def __init__(self, chessfile, append_pgn=False):
self.chessfile = chessfile
self.append_pgn = append_pgn
self.cancel = False
def initialize(self):
self.db_handle = self.chessfile.handle
self.engine = self.chessfile.engine
self.conn = self.engine.connect()
self.CHUNK = 1000
self.count_source = select([func.count()]).select_from(source)
self.ins_event = event.insert()
self.ins_site = site.insert()
self.ins_player = player.insert()
self.ins_annotator = annotator.insert()
self.ins_source = source.insert()
self.ins_game = game.insert()
self.ins_tag_game = tag_game.insert()
self.event_dict = {}
self.site_dict = {}
self.player_dict = {}
self.annotator_dict = {}
self.source_dict = {}
self.next_id = [0, 0, 0, 0, 0, 0]
self.next_id[GAME] = self.ini_names(game, GAME)
self.next_id[EVENT] = self.ini_names(event, EVENT)
self.next_id[SITE] = self.ini_names(site, SITE)
self.next_id[PLAYER] = self.ini_names(player, PLAYER)
self.next_id[ANNOTATOR] = self.ini_names(annotator, ANNOTATOR)
self.next_id[SOURCE] = self.ini_names(source, SOURCE)
def get_id(self, name, name_table, field, info=None):
if not name:
return None
orig_name = name
if field == EVENT:
name_dict = self.event_dict
name_data = self.event_data
elif field == SITE:
name_dict = self.site_dict
name_data = self.site_data
elif field == ANNOTATOR:
name_dict = self.annotator_dict
name_data = self.annotator_data
elif field == SOURCE:
name_dict = self.source_dict
name_data = self.source_data
elif field == PLAYER:
name_dict = self.player_dict
name_data = self.player_data
name = name.title().translate(removeDic)
if name in name_dict:
return name_dict[name]
if field == SOURCE:
name_data.append({'name': orig_name, 'info': info})
else:
name_data.append({'name': orig_name})
name_dict[name] = self.next_id[field]
self.next_id[field] += 1
return name_dict[name]
def ini_names(self, name_table, field):
if field != GAME and field != STAT:
s = select([name_table])
name_dict = dict([(n.name.title().translate(removeDic), n.id)
for n in self.conn.execute(s)])
if field == EVENT:
self.event_dict = name_dict
elif field == SITE:
self.site_dict = name_dict
elif field == PLAYER:
self.player_dict = name_dict
elif field == ANNOTATOR:
self.annotator_dict = name_dict
elif field == SOURCE:
self.source_dict = name_dict
s = select([func.max(name_table.c.id).label('maxid')])
maxid = self.conn.execute(s).scalar()
if maxid is None:
next_id = 1
else:
next_id = maxid + 1
return next_id
def do_cancel(self):
GLib.idle_add(self.progressbar.set_text, "")
self.cancel = True
# @profile_me
def do_import(self, filename, info=None, progressbar=None):
self.progressbar = progressbar
orig_filename = filename
count_source = self.conn.execute(self.count_source.where(source.c.name == orig_filename)).scalar()
if count_source > 0:
log.info("%s is already imported" % filename)
return
# collect new names not in they dict yet
self.event_data = []
self.site_data = []
self.player_data = []
self.annotator_data = []
self.source_data = []
# collect new games and commit them in big chunks for speed
self.game_data = []
self.tag_game_data = []
if filename.startswith("http"):
filename = download_file(filename, progressbar=progressbar)
if filename is None:
return
else:
if not os.path.isfile(filename):
log.info("Can't open %s" % filename)
return
if filename.lower().endswith(".zip") and zipfile.is_zipfile(filename):
with zipfile.ZipFile(filename, "r") as zf:
path = os.path.dirname(filename)
files = [os.path.join(path, f) for f in zf.namelist() if f.lower().endswith(".pgn")]
zf.extractall(path)
else:
files = [filename]
for pgnfile in files:
base_offset = self.chessfile.size if self.append_pgn else 0
basename = os.path.basename(pgnfile)
if progressbar is not None:
GLib.idle_add(progressbar.set_text, _("Reading %s ..." % basename))
else:
log.info("Reading %s ..." % pgnfile)
size = os.path.getsize(pgnfile)
handle = protoopen(pgnfile)
# estimated game count
all_games = max(size / 840, 1)
get_id = self.get_id
# use transaction to avoid autocommit slowness
# and to let undo importing (rollback) if self.cancel was set
trans = self.conn.begin()
try:
i = 0
for tags in read_games(handle):
if not tags:
log.info("Empty game #%s" % (i + 1))
continue
if self.cancel:
trans.rollback()
return
fenstr = tags["FEN"]
variant = tags["Variant"]
if variant:
if "fischer" in variant.lower() or "960" in variant:
variant = "Fischerandom"
else:
variant = variant.lower().capitalize()
# Fixes for some non statndard Chess960 .pgn
if fenstr and variant == "Fischerandom":
parts = fenstr.split()
parts[0] = parts[0].replace(".", "/").replace("0", "")
if len(parts) == 1:
parts.append("w")
parts.append("-")
parts.append("-")
fenstr = " ".join(parts)
if variant:
if variant not in name2variant:
log.info("Unknown variant: %s" % variant)
continue
variant = name2variant[variant].variant
if variant == NORMALCHESS:
# lichess uses tag [Variant "Standard"]
variant = 0
else:
variant = 0
if basename == "eco.pgn":
white = tags["Opening"]
black = tags["Variation"]
else:
white = tags["White"]
black = tags["Black"]
event_id = get_id(tags["Event"], event, EVENT)
site_id = get_id(tags["Site"], site, SITE)
date = tags["Date"]
game_round = tags['Round']
white_id = get_id(white, player, PLAYER)
black_id = get_id(black, player, PLAYER)
result = tags["Result"]
if result in pgn2Const:
result = pgn2Const[result]
else:
result = RUNNING
white_elo = tags['WhiteElo']
black_elo = tags['BlackElo']
time_control = tags["TimeControl"]
eco = tags["ECO"][:3]
fen = tags["FEN"]
board_tag = int(tags["Board"]) if "Board" in tags else 0
annotator_id = get_id(tags["Annotator"], annotator, ANNOTATOR)
source_id = get_id(orig_filename, source, SOURCE, info=info)
ply_count = tags["PlyCount"] if "PlyCount" in tags else 0
offset = base_offset + int(tags["offset"])
self.game_data.append({
'offset': offset,
'offset8': (offset >> 3) << 3,
'event_id': event_id,
'site_id': site_id,
'date': date,
'round': game_round,
'white_id': white_id,
'black_id': black_id,
'result': result,
'white_elo': white_elo,
'black_elo': black_elo,
'ply_count': ply_count,
'eco': eco,
'fen': fen,
'variant': variant,
'board': board_tag,
'time_control': time_control,
'annotator_id': annotator_id,
'source_id': source_id,
})
for tag in tags:
if tag not in dedicated_tags and tag not in other_game_tags and tags[tag]:
self.tag_game_data.append({
'game_id': self.next_id[GAME],
'tag_name': tag,
'tag_value': tags[tag],
})
self.next_id[GAME] += 1
i += 1
if len(self.game_data) >= self.CHUNK:
if self.event_data:
self.conn.execute(self.ins_event, self.event_data)
self.event_data = []
if self.site_data:
self.conn.execute(self.ins_site, self.site_data)
self.site_data = []
if self.player_data:
self.conn.execute(self.ins_player,
self.player_data)
self.player_data = []
if self.annotator_data:
self.conn.execute(self.ins_annotator,
self.annotator_data)
self.annotator_data = []
if self.source_data:
self.conn.execute(self.ins_source, self.source_data)
self.source_data = []
if self.tag_game_data:
self.conn.execute(self.ins_tag_game, self.tag_game_data)
self.tag_game_data = []
self.conn.execute(self.ins_game, self.game_data)
self.game_data = []
if progressbar is not None:
GLib.idle_add(progressbar.set_fraction, i / float(all_games))
GLib.idle_add(progressbar.set_text, _(
"%(counter)s game headers from %(filename)s imported" % ({"counter": i, "filename": basename})))
else:
log.info("From %s imported %s" % (pgnfile, i))
if self.event_data:
self.conn.execute(self.ins_event, self.event_data)
self.event_data = []
if self.site_data:
self.conn.execute(self.ins_site, self.site_data)
self.site_data = []
if self.player_data:
self.conn.execute(self.ins_player, self.player_data)
self.player_data = []
if self.annotator_data:
self.conn.execute(self.ins_annotator, self.annotator_data)
self.annotator_data = []
if self.source_data:
self.conn.execute(self.ins_source, self.source_data)
self.source_data = []
if self.tag_game_data:
self.conn.execute(self.ins_tag_game, self.tag_game_data)
self.tag_game_data = []
if self.game_data:
self.conn.execute(self.ins_game, self.game_data)
self.game_data = []
if progressbar is not None:
GLib.idle_add(progressbar.set_fraction, i / float(all_games))
GLib.idle_add(progressbar.set_text, _(
"%(counter)s game headers from %(filename)s imported" % ({"counter": i, "filename": basename})))
else:
log.info("From %s imported %s" % (pgnfile, i))
trans.commit()
if self.append_pgn:
# reopen database to write
self.db_handle.close()
with protosave(self.chessfile.path, self.append_pgn) as self.db_handle:
log.info("Append from %s to %s" % (pgnfile, self.chessfile.path))
handle.seek(0)
self.db_handle.writelines(handle)
handle.close()
if self.chessfile.scoutfish is not None:
# create new .scout from pgnfile we are importing
from pychess.Savers.pgn import scoutfish_path
args = [scoutfish_path, "make", pgnfile, "%s" % base_offset]
output = subprocess.check_output(args, stderr=subprocess.STDOUT).decode()
# append it to our existing one
if output.find("Processing...done") > 0:
old_scout = self.chessfile.scoutfish.db
new_scout = os.path.splitext(pgnfile)[0] + '.scout'
with open(old_scout, "ab") as file1, open(new_scout, "rb") as file2:
file1.write(file2.read())
self.chessfile.handle = protoopen(self.chessfile.path)
except SQLAlchemyError as e:
trans.rollback()
log.info("Importing %s failed! \n%s" % (pgnfile, e))
def read_games(handle):
"""Based on chess.pgn.scan_headers() from Niklas Fiekas python-chess"""
in_comment = False
game_headers = None
game_pos = None
last_pos = 0
line = handle.readline()
# scoutfish creates game offsets at previous game end
line_end_fix = 2 if line.endswith("\r\n") else 1
while line:
# Skip single line comments.
if line.startswith("%"):
last_pos += len(line)
line = handle.readline()
continue
# Reading a header tag. Parse it and add it to the current headers.
if not in_comment and line.startswith("["):
tag_match = TAG_REGEX.match(line)
if tag_match:
if game_pos is None:
game_headers = collections.defaultdict(str)
game_pos = last_pos
tag_value = tag_match.group(2)
tag_value = tag_value.replace("\\\"", "\"")
tag_value = tag_value.replace("\\\\", "\\")
if handle.pgn_encoding != PGN_ENCODING:
tag_value = tag_value.encode(PGN_ENCODING).decode(handle.pgn_encoding)
game_headers[tag_match.group(1)] = tag_value
last_pos += len(line)
line = handle.readline()
continue
# Reading movetext. Update parser state in_comment in order to skip
# comments that look like header tags.
if (not in_comment and "{" in line) or (in_comment and "}" in line):
in_comment = line.rfind("{") > line.rfind("}")
# Reading movetext. If there were headers, previously, those are now
# complete and can be yielded.
if game_pos is not None:
game_headers["offset"] = max(0, game_pos - line_end_fix)
yield game_headers
game_pos = None
last_pos += len(line)
line = handle.readline()
# Yield the headers of the last game.
if game_pos is not None:
game_headers["offset"] = max(0, game_pos - line_end_fix)
yield game_headers
