gostyle: the basic library, intitial commit.
[gostyle.git] / utils / godb_models.py
blob69eeee767a0b150f8c9d2dd1318ac819e081849d
1 import os
2 from itertools import chain
4 import logging
5 import re
7 from sqlalchemy import Table, Column, Integer, ForeignKey, Text, Date, Float, Enum, UniqueConstraint, PickleType
8 from sqlalchemy.orm import relationship, backref
9 from sqlalchemy.ext.declarative import declarative_base
11 import utils
12 from rank import Rank
13 from colors import *
15 """
16 This contains the sqlalchemy ORM models which also form basic of our datastructures, have a look at it.
18 """
20 Base = declarative_base()
22 class ProcessingError(Exception):
23 pass
25 class SchizophrenicPlayerError(Exception):
26 """Used in context of problems with games between the same players.
27 E.g. Anonymous vs. Anonymous"""
28 pass
32 ## hack to workaround this bug: http://bugs.python.org/issue5876
33 ## > Oh ok, gotcha: repr() always returns a str string. If obj.__repr__() returns a
34 ## > Unicode string, the string is encoded to the default encoding. By default, the
35 ## > default encoding is ASCII.
36 ## => unicode chars in repr cause "ordinal not in range err"
37 import functools
38 import misc
39 def ununicode(f):
40 @functools.wraps(f)
41 def g(*args, **kwargs):
42 return misc.unicode2ascii(f(*args, **kwargs))
43 return g
45 class Player(Base):
46 """Class (and ORM Table) about go player.
47 The name must be unique (born name, ..). This class should have one instance (db record)
48 for one go player.
50 Player may change name, rank, .. in time, or use different nicknames, etc.
51 The consistency (so that all these variations are connected) is maintained
52 together with the PlayerInTime.
53 """
54 __tablename__ = 'player'
55 id = Column(Integer, primary_key=True, index=True)
56 name = Column(Text, nullable=False, unique=True, index=True)
57 note = Column(Text)
58 # list in_times from backrefs
60 #__table_args__ = ( UniqueConstraint('name'), )
62 def __init__(self, name, note=u''):
63 self.name = name
64 self.note = note
66 def iter_games_as(self, color, pit_filter=lambda pit:True):
67 return chain.from_iterable(pit.iter_games_as(color)
68 for pit in self.in_times if pit_filter(pit) )
70 def iter_one_side_associations(self,
71 pit_filter=lambda pit:True,
72 **kwargs):
73 return chain.from_iterable( pit.iter_one_side_associations(**kwargs)
74 for pit in self.in_times if pit_filter(pit) )
75 # Shortcuts
76 def iter_games_as_white(self, **kwargs):
77 return self.iter_games_as(PLAYER_COLOR_WHITE, **kwargs)
78 def iter_games_as_black(self, **kwargs):
79 return self.iter_games_as(PLAYER_COLOR_BLACK, **kwargs)
80 def iter_games(self):
81 return chain(self.iter_games_as_black(), self.iter_games_as_white())
83 def __str__(self):
84 return self.name
86 @ununicode
87 def __repr__(self):
88 return u"Player(%s, '%s','%s')" % (self.id,
89 self.name,
90 self.note)
92 import pickle
94 class PlayerInTime(Base):
95 """Captures evolution of players in time - change of rank, name, different identities."""
96 __tablename__ = 'player_in_time'
97 id = Column(Integer, primary_key=True, index=True)
99 player_id = Column(Integer, ForeignKey('player.id'), index=True)
100 player = relationship("Player", backref=backref('in_times', order_by=id))
102 name = Column(Text)
104 rank = Column(PickleType) # (pickler=pickle))
105 note = Column(Text)
106 # list games_as_black from backrefs
107 # list games_as_white from backrefs
109 def __init__(self, player, name='', rank=None, note=''):
110 if isinstance(rank, basestring):
111 rank = Rank.from_string(rank)
113 self.player = player
114 self.name = name
115 self.rank = rank
116 self.note = note
118 def get_games_as(self, color):
119 if color == PLAYER_COLOR_BLACK :
120 return self.games_as_black
121 if color == PLAYER_COLOR_WHITE :
122 return self.games_as_white
123 raise KeyError(color)
125 def iter_games_as(self, color):
126 return iter(self.get_games_as(color))
128 def iter_one_side_associations(self,
129 color_filter=lambda color:True,
130 game_filter=lambda game:True ):
131 return ( OneSideListAssociation(game, color)
132 for color in PLAYER_COLORS if color_filter(color)
133 for game in self.iter_games_as(color) if game_filter(game) )
135 def __str__(self):
136 return self.name + ( " (%s)"%(self.rank) if self.rank else '')
138 def str2(self):
139 return self.name + ( " [%s]"%(self.rank) if self.rank else '')
141 @ununicode
142 def __repr__(self):
143 return u"PlayerInTime(%s, %s, '%s', '%s', '%s')" % (
144 self.id,
145 repr(self.player),
146 self.name,
147 self.rank,
148 self.note )
150 class Game(Base):
151 """Class (and ORM Table) holding game information like
152 - sgf filename
153 - info about players - who played black, who played white
154 - sgf header with further info
156 __tablename__ = 'game'
157 id = Column(Integer, primary_key=True, index=True)
158 sgf_file = Column(Text, nullable=False)
160 black_id = Column(Integer, ForeignKey('player_in_time.id'), index=True)
161 white_id = Column(Integer, ForeignKey('player_in_time.id'), index=True)
163 black = relationship("PlayerInTime", primaryjoin="PlayerInTime.id==Game.black_id",
164 backref=backref('games_as_black', order_by=id))
165 white = relationship("PlayerInTime", primaryjoin="PlayerInTime.id==Game.white_id",
166 backref=backref('games_as_white', order_by=id))
168 sgf_header = Column(PickleType)
170 # We store the whole header instead of these
172 #date = Column(Text)
173 #komi = Column(Float)
174 #handicap = Column(Integer)
175 #size = Column(Integer)
176 #result = Column(Text)
177 #note = Column(Text)
179 def __init__(self, sgf_file, black, white, sgf_header={}):
180 self.sgf_file = sgf_file
181 self.black = black
182 self.white = white
183 self.sgf_header = sgf_header
185 @ununicode
186 def __repr__(self):
187 return u"Game(%s, '%s', '%s', '%s')" %(
188 self.id,
189 self.sgf_file,
190 repr(self.white) if self.white else '',
191 repr(self.black) if self.black else '')
193 def abs_path(self):
194 return os.path.abspath(self.sgf_file)
196 def iter_pit_color(self):
197 yield (self.black, PLAYER_COLOR_BLACK)
198 yield (self.white, PLAYER_COLOR_WHITE)
200 def get_player_by_color(self, color):
201 for gpit, gcolor in self.iter_pit_color():
202 if color == gcolor:
203 return gpit.player
204 raise ValueError("Wrong color '%s'."%color)
206 def get_player_color(self, player):
207 if self.black.player == self.white.player :
208 # we cannot expect for this method to return different values for one player...
209 # (so this would always return black, because it has no way of knowing if we ask for
210 # black or white player)
211 raise SchizophrenicPlayerError("Asked for color for game between identical players: %s"%(self,))
213 for gpit, gcolor in self.iter_pit_color():
214 if player == gpit.player:
215 return gcolor
217 raise ValueError("Game %s is not a game of %s."%(repr(self), repr(player)))
219 def get_year(self, try_filename_prefix=True):
220 # Year from DT field of sgf file
221 dt = self.sgf_header.get('DT', 'Unknown')
222 year = utils.get_year(dt)
224 # try to guess name from filename prefix (e.g. gogod)
225 if year == None and try_filename_prefix:
226 fn = os.path.basename(self.sgf_file)[:4]
227 return utils.get_year(fn)
229 # return year or None if failure
230 return year
232 def open_in_viewer(self):
233 utils.viewer_open(self.abs_path())
235 game_list_association = Table('game_list_association', Base.metadata,
236 Column('game_list_id', Integer, ForeignKey('game_list.id'), index=True),
237 Column('game_id', Integer, ForeignKey('game.id'), index=True)
240 class GameList(Base):
241 """List of games.
243 __tablename__ = 'game_list'
244 id = Column(Integer, primary_key=True, index=True)
245 name = Column(Text, nullable=False, unique=True, index=True)
247 games = relationship('Game', secondary=game_list_association, backref='game_lists')
249 def __init__(self, name, games=None):
250 self.name = name
251 if games != None:
252 assert not self.games
253 self.games = list(games)
255 def iter_players_black(self):
256 """Iterate players who played in a game (from this list) as black."""
257 for game in self.games:
258 yield game.black.player
260 def iter_players_white(self):
261 """Look at self.get_players_black and guess."""
262 for game in self.games:
263 yield game.white.player
265 def iter_players(self):
266 """Iterate players who played a game from this list."""
267 return chain(self.iter_players_black(), self.iter_players_white())
269 def append(self, game):
270 self.games.append(game)
272 #def __str__(self):
273 # ret = [ self.name ] + map(str, self.games)
275 # return '\n'.join(ret)
277 def __getitem__(self, val):
278 return self.games[val]
280 def __len__(self):
281 return len(self.games)
283 @ununicode
284 def __repr__(self):
285 return "GameList(%s, '%s', #games = %d)" %( self.id, self.name, len(self) )
288 class Merger:
289 def __init__(self):
290 pass
291 def __repr__(self):
292 return self.__class__.__name__ + "()"
293 def start(self, bw_gen):
294 raise NotImplementedError
295 def add(self, result, color):
296 raise NotImplementedError
297 def finish(self):
298 raise NotImplementedError
301 class OneSideListAssociation(Base):
302 __tablename__ = 'one_side_list_association'
303 id = Column(Integer, primary_key=True, index=True)
304 one_side_list_id = Column(Integer, ForeignKey('one_side_list.id'), index=True)
305 game_id = Column(Integer, ForeignKey('game.id'), index=True)
307 # what is the color of the player of interest in this game?
308 color = Column(Enum(PLAYER_COLOR_BLACK, PLAYER_COLOR_WHITE))
309 game = relationship("Game", backref="one_side_lists_assoc")
311 # one game ( for one side ) can be in one game list only once
312 __table_args__ = ( UniqueConstraint('one_side_list_id', 'game_id', 'color'), )
314 def __init__(self, game, color):
315 self.game = game
316 self.color = color
318 def __iter__(self):
319 yield self.game
320 yield self.color
322 @ununicode
323 def __repr__(self):
324 return u"OneSideListAssociation(%s, '%s')" %( repr(self.game), self.color )
326 class OneSideList(Base):
327 """List of games, for e.g. players with 10kyu, Honinbo Shusaku's games, ...
329 Note that the list distinguishes between sides. That is, if you are interested
330 in both sides (default behaviour of the `add` method), the game will be added
331 twice - once for black, once for white.
333 __tablename__ = 'one_side_list'
334 id = Column(Integer, primary_key=True, index=True)
335 name = Column(Text, nullable=False, unique=True, index=True)
337 list_assocs = relationship('OneSideListAssociation', backref='one_side_list')
339 def __init__(self, name, assocs=None):
340 self.name = name
341 if assocs != None:
342 assert not self.list_assocs
343 self.list_assocs = list(assocs)
345 def __getitem__(self, val):
346 return self.list_assocs[val]
348 def batch_add(self, games, color):
349 """Add games played with one color in batch."""
350 self.list_assocs += [ OneSideListAssociation(game, color) for game in games ]
352 def add(self, game, player=None, color=None):
353 """Adds game to the list. If @player (or @color) specified, adds only
354 one side of the game - the one that @player played (or played with @color).
355 Otw. both sides get added (game is added twice - once for black, once for white)
357 if player != None:
358 pc = game.get_player_color(player)
359 if color and color != cp:
360 raise ValueError( """Provided color (%s) is different from provided player's (%s) color in the game %s."""%
361 ( color, player, game ))
362 color = pc
364 if color != None:
365 # if color of the desired player specified
366 citer = ( color, )
367 else:
368 # add both black's game and white's game
369 citer = PLAYER_COLORS
371 for color in citer:
372 self.list_assocs.append(OneSideListAssociation(game, color))
374 def for_one_side_list(self, merger, bw_processor):
376 Processes the whole OneSideList, so that @bw_processor is called on every game. And the
377 result of interest (black or white) is added to @merger, via @merger.add(result, color).
378 At the end @merger.finish() is called and this should return the desired data.
380 #assert isinstance(merger, Merger)
382 merger.start(bw_processor)
384 for ga in self.list_assocs:
385 try:
386 black, white = bw_processor(ga.game)
387 except ProcessingError as exc:
388 logging.debug("Exception %s occured in processing the game %s, skipping!!"%(repr(exc), ga.game))
389 continue
390 except Exception as exc:
391 logging.exception("Exception %s occured in processing the game %s!!"%(repr(exc), ga.game))
392 raise
393 #continue
395 desired = black if ga.color == PLAYER_COLOR_BLACK else white
396 merger.add(desired, ga.color)
398 return merger.finish()
400 def __len__(self):
401 return len(self.list_assocs)
403 def __str__(self):
404 ret = [ self.name ]
405 for ga in self.list_assocs:
406 ret.append("%s : %s"%(ga.color, ga.game))
408 return '\n'.join(ret)
410 @ununicode
411 def __repr__(self):
412 return "OneSideList(%s, '%s', #games = %d)"%( self.id, self.name, len(self) )
414 class DataMap(Base):
416 One DataMap holds info about the mapping:
417 OneSideList -> ImageData
419 __tablename__ = 'datamap'
420 id = Column(Integer, primary_key=True, index=True)
421 name = Column(Text, nullable=False, unique=True, index=True)
423 # information about the image domain
424 image_types = Column(PickleType)
425 image_annotations = Column(PickleType)
427 relations = relationship("DataMapRelation", backref='datamap')
429 def add(self, one_side_list, image):
430 self.relations.append(DataMapRelation(one_side_list=one_side_list,
431 image=image))
432 def __len__(self):
433 return len(self.relations)
435 def __getitem__(self, val):
436 return self.relations[val]
438 @ununicode
439 def __repr__(self):
440 return "DataMap(%d, '%s', #relations = %d )"%(self.id, self.name, len( self.relations))
442 class DataMapRelation(Base):
444 One OneSideList gets mapped to data (usually a python vector).
446 __tablename__ = 'datamap_relation'
447 id = Column(Integer, primary_key=True, index=True)
448 # id of the current dataset
449 datamap_id = Column(Integer, ForeignKey('datamap.id'), index=True)
450 # domain
451 one_side_list_id = Column(Integer, ForeignKey('one_side_list.id'), index=True)
452 # image
453 image_id = Column(Integer, ForeignKey('image_data.id'), index=True)
455 one_side_list = relationship("OneSideList")#, backref='relations')
456 image = relationship("ImageData")
458 def __iter__(self):
459 yield self.one_side_list
460 yield self.image
462 def __repr__(self):
463 return "DataMapRelation(%s,%s)" % (repr(self.one_side_list),
464 repr(self.image))
467 class ImageData(Base):
468 """ Class used to hold python-pickled data under unique name. Meant to be
469 used for holding right side of the mapping defined by DataMapRelation,
470 so that multiple OneSideLists may share the same image.
472 __tablename__ = 'image_data'
473 id = Column(Integer, primary_key=True, index=True)
474 # e.g. 'style: Otake Hideo'
475 name = Column(Text, nullable=False, unique=True, index=True)
476 # e.g. the style vector itself
477 data = Column(PickleType)
479 def __init__(self, name, data):
480 self.name = name
481 self.data = data
483 @ununicode
484 def __repr__(self):
485 return "ImageData(%s, %s, %s)"%(self.id, self.name, self.data)