utils/godb_session.py

   1 from sqlalchemy import create_engine, and_
   2 from sqlalchemy.orm import sessionmaker, aliased
   3 from sqlalchemy.orm.session import Session
   4
   5 import logging
   6 import os
   7 import re
   8
   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
  13
  14 from sgf_load import load_sgf_file_headers, ParseError
  15
  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)
  31
  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
  39
  40         pits = self.query(PlayerInTime).filter( PlayerInTime.player == player )
  41
  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 )
  48
  49         pit_all = pits.all()
  50         if len(pit_all):
  51             return pit_all[0]
  52
  53         return PlayerInTime(player, current_name, current_rank, current_note)
  54
  55     def godb_sgf_to_game(self, filename):
  56         """
  57         Creates a Game object from .sgf file.
  58
  59         Currently, only sgf files with a single gametree are supported.
  60
  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
  69
  70         if not headers:
  71             logging.warn("No headers in .sgf file '%s', skipping."%(filename,))
  72             return None
  73
  74         if len(headers) > 1:
  75             logging.warn("More game trees in a file, skipping '%s'."%(filename,))
  76             return None
  77
  78         hd = headers[0]
  79
  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)
  88
  89         return Game( filename.decode('utf-8'), pb, pw, hd )
  90
  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)
  94
  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).
  99
 100         Both players in each of the games scanned are added into the session.
 101         (see self.godb_sgf_to_game)
 102
 103         The gamelist is returned and NOT added into the session.
 104         """
 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))
 111
 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()
 118
 119         t.stop_n_log('  Total time', 'Game')
 120
 121         if gamelist == None:
 122             gamelist = GameList("Games from '%s'."%(directory,))
 123
 124         gamelist.games += games
 125         logging.info("Added %d games to: %s"%(len(games), gamelist))
 126
 127         return gamelist
 128
 129     """
 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()
 142     """
 143
 144 _godb_session = sessionmaker(class_=GodbSession)
 145
 146 def godb_session_maker(filename, echo=False):
 147     engine = create_engine('sqlite:///%s'%filename, echo=echo)
 148     Base.metadata.create_all(engine)
 149
 150     s = _godb_session(bind=engine)
 151     # for wingide completion...
 152     isinstance(s, GodbSession)
 153     return s
 154
 155 if __name__ == '__main__':
 156     import logging
 157     logging.getLogger().setLevel(logging.DEBUG)
 158
 159     s = godb_session_maker(filename=':memory:')#, echo=True)
 160
 161     g = s.godb_sgf_to_game('../TEST_FILES/games/1990-00-00b.sgf')
 162     print unicode(g)
 163
 164     gl = s.godb_scan_dir_as_gamelist('../TEST_FILES/games')
 165
 166     s.add(gl)
 167     s.commit()
 168
 169     print gl
 170     for g in gl:
 171         print "%20s vs. %s" % (g.white,  g.black)
 172
 173