gostyle: the basic library, intitial commit.
[gostyle.git] / utils / godb_session.py
blob378718feecdf83ecfe740a5459e4de8f2ba1f672
1 from sqlalchemy import create_engine, and_
2 from sqlalchemy.orm import sessionmaker, aliased
3 from sqlalchemy.orm.session import Session
5 import logging
6 import os
7 import re
9 from godb_models import Player, PlayerInTime, Game, GameList, OneSideList, OneSideListAssociation, DataMap, DataMapRelation, ImageData, Base
10 from rank import Rank
11 import misc
12 import timer
14 from sgf_load import load_sgf_file_headers, ParseError
16 class GodbSession(Session):
17 def godb_get_player(self, name, note=u''):
18 """Looks if the player with @name is in the DB and returns it.
19 Otw. creates a new player with these attributes.
20 This new player is NOT added into the session.
21 """
22 pls = self.query(Player).filter(Player.name==name).all()
23 assert len(pls) <= 1
24 if len(pls) == 1:
25 player = pls[0]
26 if player.note != note:
27 logging.warn("%s has different note than '%s'"%(repr(player), note))
28 return player
29 if len(pls) == 0:
30 return Player(name, note)
32 def godb_get_player_in_time(self, name, current_name=None, current_rank=None, current_note=''):
33 """
34 NOT adding anything into the session.
35 """
36 player = self.godb_get_player(name)
37 if current_name == None:
38 current_name = name
40 pits = self.query(PlayerInTime).filter( PlayerInTime.player == player )
42 if current_name:
43 pits = pits.filter( PlayerInTime.name == current_name )
44 if current_rank:
45 pits = pits.filter( PlayerInTime.rank == current_rank )
46 if current_note:
47 pits = pits.filter( PlayerInTime.note == current_note )
49 pit_all = pits.all()
50 if len(pit_all):
51 return pit_all[0]
53 return PlayerInTime(player, current_name, current_rank, current_note)
55 def godb_sgf_to_game(self, filename):
56 """
57 Creates a Game object from .sgf file.
59 Currently, only sgf files with a single gametree are supported.
61 Does NOT add the game in the session but it
62 DOES ADD players in the game in there.
63 """
64 try:
65 headers = load_sgf_file_headers(filename)
66 except ParseError:
67 logging.warn("Could not parse '%s', probably not a .sgf file, skipping."%(filename,))
68 return None
70 if not headers:
71 logging.warn("No headers in .sgf file '%s', skipping."%(filename,))
72 return None
74 if len(headers) > 1:
75 logging.warn("More game trees in a file, skipping '%s'."%(filename,))
76 return None
78 hd = headers[0]
80 # load players' names and ranks
81 # we add them to the session to have consistency and correctly interjoined objects
82 # (e.g. when pw == pb (anonymous) then only the first call actually
83 # creates a new object. The second call uses the same object.
84 pw = self.godb_get_player_in_time(hd.get('PW', ''), current_rank=Rank.from_string(hd.get('WR','')))
85 self.add(pw)
86 pb = self.godb_get_player_in_time(hd.get('PB', ''), current_rank=Rank.from_string(hd.get('BR','')))
87 self.add(pb)
89 return Game( filename.decode('utf-8'), pb, pw, hd )
91 def godb_add_dir_as_gamelist(self, *args, **kwargs):
92 logging.warn("deprecated call to godb_add_dir_as_gamelist")
93 return self.godb_scan_dir_as_gamelist(self, *args, **kwargs)
95 def godb_scan_dir_as_gamelist(self, directory, gamelist=None):
96 """Recursively scans the @directory for sgf files.
97 The valid games are added into a gamelist (either provided by @gamelist kwarg,
98 or new if @gamelist == None).
100 Both players in each of the games scanned are added into the session.
101 (see self.godb_sgf_to_game)
103 The gamelist is returned and NOT added into the session.
105 t = timer.Timer()
106 games = []
107 t.start()
108 for filepath in misc.iter_files(directory):
109 if re.search('sgf$', filepath):
110 logging.debug("Scanning '%s'"%(filepath))
112 # create Game object from the sgf file
113 t.start()
114 game = self.godb_sgf_to_game(filepath)
115 if game:
116 games.append(game)
117 t.stop()
119 t.stop_n_log(' Total time', 'Game')
121 if gamelist == None:
122 gamelist = GameList("Games from '%s'."%(directory,))
124 gamelist.games += games
125 logging.info("Added %d games to: %s"%(len(games), gamelist))
127 return gamelist
130 ## TODO make it faster!!!
131 def godb_list_player_games_white(self, pits):
132 #pits = self.query(PlayerInTime.id).filter(PlayerInTime.player_id == player.id).all()
133 #pits = player.in_times
134 return self.query(Game).filter(Game.white_id.in_(pits) ).all()
135 def godb_list_player_games_black(self, pits):
136 #pits = ( pit.id for pit in player.in_times )
137 #pits = self.query(PlayerInTime.id).filter(PlayerInTime.player_id == player.id).all()
138 return self.query(Game).filter(Game.black_id.in_(pits) ).all()
139 #return self.query(Game).\
140 # join(PlayerInTime, Game.black_id==PlayerInTime.id).\
141 # filter(PlayerInTime.player_id == player.id).all()
144 _godb_session = sessionmaker(class_=GodbSession)
146 def godb_session_maker(filename, echo=False):
147 engine = create_engine('sqlite:///%s'%filename, echo=echo)
148 Base.metadata.create_all(engine)
150 s = _godb_session(bind=engine)
151 # for wingide completion...
152 isinstance(s, GodbSession)
153 return s
155 if __name__ == '__main__':
156 import logging
157 logging.getLogger().setLevel(logging.DEBUG)
159 s = godb_session_maker(filename=':memory:')#, echo=True)
161 g = s.godb_sgf_to_game('../TEST_FILES/games/1990-00-00b.sgf')
162 print unicode(g)
164 gl = s.godb_scan_dir_as_gamelist('../TEST_FILES/games')
166 s.add(gl)
167 s.commit()
169 print gl
170 for g in gl:
171 print "%20s vs. %s" % (g.white, g.black)