soltesz.py

   1 import os
   2 import sys
   3 import pickle
   4 import inspect
   5 import shutil
   6 from config import config
   7 config = config()
   8
   9 DEBUG= 0
  10 PICKLE_PATH="pdb"
  11
  12 def dbLoad(name):
  13         return SPickle().load(name)
  14
  15 def dbExists(name):
  16         #if self.config.debug:
  17         #       name = "debug.%s" % name
  18         return SPickle().exists(name)
  19
  20 def dbDump(name, obj=None):
  21         # depth of the dump is 2 now, since we're redirecting to '.dump'
  22         return SPickle().dump(name, obj, 2)
  23
  24 def if_cached_else_refresh(cond, refresh, name, function):
  25         s = SPickle()
  26         if refresh:
  27                 if not config.debug and s.exists("production.%s" % name):
  28                         s.remove("production.%s" % name)
  29                 if config.debug and s.exists("debug.%s" % name):
  30                         s.remove("debug.%s" % name)
  31
  32         return if_cached_else(cond, name, function)
  33
  34 def if_cached_else(cond, name, function):
  35         s = SPickle()
  36         if (cond and s.exists("production.%s" % name)) or \
  37            (cond and config.debug and s.exists("debug.%s" % name)):
  38                 o = s.load(name)
  39         else:
  40                 o = function()
  41                 if cond:
  42                         s.dump(name, o) # cache the object using 'name'
  43         return o
  44
  45 class SPickle:
  46         def __init__(self):
  47                 pass
  48
  49         def if_cached_else(self, cond, name, function):
  50                 if cond and self.exists("production.%s" % name):
  51                         o = self.load(name)
  52                 else:
  53                         o = function()
  54                         if cond:
  55                                 self.dump(name, o)      # cache the object using 'name'
  56                 return o
  57
  58         def __file(self, name):
  59                 return "%s/%s.pkl" % (PICKLE_PATH, name)
  60
  61         def exists(self, name):
  62                 return os.path.exists(self.__file(name))
  63
  64         def remove(self, name):
  65                 return os.remove(self.__file(name))
  66
  67         def load(self, name):
  68                 """
  69                 In debug mode, we should fail if neither file exists.
  70                         if the debug file exists, reset name
  71                         elif the original file exists, make a copy, reset name
  72                         else neither exist, raise an error
  73                 Otherwise, it's normal mode, if the file doesn't exist, raise error
  74                 Load the file
  75                 """
  76
  77                 if config.debug:
  78                         if self.exists("debug.%s" % name):
  79                                 name = "debug.%s" % name
  80                         elif self.exists("production.%s" % name):
  81                                 debugname = "debug.%s" % name
  82                                 if not self.exists(debugname):
  83                                         name = "production.%s" % name
  84                                         shutil.copyfile(self.__file(name), self.__file(debugname))
  85                                 name = debugname
  86                         else:   # neither exist
  87                                 raise Exception, "No such pickle based on %s" % self.__file("debug.%s" % name)
  88                 else:
  89                         if not self.exists("production.%s" % name):
  90                                 raise Exception, "No such file %s" % name
  91                         name = "production.%s" % name
  92
  93                 print "loading %s" % self.__file(name)
  94                 f = open(self.__file(name), 'r')
  95                 o = pickle.load(f)
  96                 f.close()
  97                 return o
  98
  99
 100         # use the environment to extract the data associated with the local
 101         # variable 'name'
 102         def dump(self, name, obj=None, depth=1):
 103                 if obj == None:
 104                         o = inspect.getouterframes(inspect.currentframe())
 105                         up1 = o[depth][0] # get the frame one prior to (up from) this frame
 106                         argvals = inspect.getargvalues(up1)
 107                         # TODO: check that 'name' is a local variable; otherwise this would fail.
 108                         obj = argvals[3][name] # extract the local variable name 'name'
 109                 if not os.path.isdir("%s/" % PICKLE_PATH):
 110                         os.mkdir("%s" % PICKLE_PATH)
 111                 if config.debug:
 112                         name = "debug.%s" % name
 113                 else:
 114                         name = "production.%s" % name
 115                 f = open(self.__file(name), 'w')
 116                 pickle.dump(obj, f)
 117                 f.close()
 118                 return
 119
 120
 121 COMMAND_TIMEOUT = 60
 122 ssh_options = { 'StrictHostKeyChecking':'no',
 123                                 'BatchMode':'yes',
 124                                 'PasswordAuthentication':'no',
 125                                 'ConnectTimeout':'%s' % COMMAND_TIMEOUT}
 126 from select import select
 127 class CMD:
 128         def __init__(self):
 129                 pass
 130
 131         def run_noexcept(self, cmd):
 132
 133                 (f_in, f_out, f_err) = os.popen3(cmd)
 134                 lout, lin, lerr = select([f_out,f_err], [], [], COMMAND_TIMEOUT*2)
 135                 if len(lin) == 0 and len(lout) == 0 and len(lerr) == 0:
 136                         # Reached a timeout!
 137                         print "TODO: kill subprocess: '%s'" % cmd
 138                         # TODO: kill subprocess??
 139                         return ("", "SCRIPTTIMEOUT")
 140                 o_value = f_out.read()
 141                 e_value = ""
 142                 if o_value == "":       # An error has occured
 143                         e_value = f_err.read()
 144
 145                 o_value = o_value.strip()
 146                 e_value = e_value.strip()
 147
 148                 f_out.close()
 149                 f_in.close()
 150                 f_err.close()
 151                 return (o_value, e_value)
 152
 153         def run(self, cmd):
 154
 155                 (f_in, f_out, f_err) = os.popen3(cmd)
 156                 value = f_out.read()
 157                 if value == "":
 158                         raise Exception, f_err.read()
 159                 value = value.strip()
 160
 161                 f_out.close()
 162                 f_in.close()
 163                 f_err.close()
 164                 return value
 165
 166
 167
 168 class SSH(CMD):
 169         def __init__(self, user, host, options = ssh_options):
 170                 self.options = options
 171                 self.user = user
 172                 self.host = host
 173                 return
 174
 175         def __options_to_str(self):
 176                 options = ""
 177                 for o,v in self.options.iteritems():
 178                         options = options + "-o %s=%s " % (o,v)
 179                 return options
 180
 181         def run(self, cmd):
 182                 cmd = "ssh %s %s@%s '%s'" % (self.__options_to_str(),
 183                                                                         self.user, self.host, cmd)
 184                 return CMD.run(self, cmd)
 185
 186         def get_file(self, rmt_filename, local_filename=None):
 187                 if local_filename == None:
 188                         local_filename = "./"
 189                 cmd = "scp -B %s %s@%s:%s %s" % (self.__options_to_str(),
 190                                                                         self.user, self.host,
 191                                                                         rmt_filename, local_filename)
 192                 # output :
 193                 #       errors will be on stderr,
 194                 #   success will have a blank stderr...
 195                 return CMD.run_noexcept(self, cmd)
 196
 197         def run_noexcept(self, cmd):
 198                 cmd = "ssh %s %s@%s '%s'" % (self.__options_to_str(),
 199                                                                         self.user, self.host, cmd)
 200                 return CMD.run_noexcept(self, cmd)
 201
 202         def runE(self, cmd):
 203                 cmd = "ssh %s %s@%s '%s'" % (self.__options_to_str(),
 204                                                                         self.user, self.host, cmd)
 205                 if ( DEBUG == 1 ):
 206                         print cmd,
 207                 (f_in, f_out, f_err) = os.popen3(cmd)
 208
 209                 value = f_out.read()
 210                 if value == "": # An error has occured
 211                         value = f_err.read()
 212                         value = value.strip()
 213
 214                 if ( DEBUG == 1 ):
 215                         print " == %s" % value
 216                 f_out.close()
 217                 f_in.close()
 218                 f_err.close()
 219                 return value.strip()
 220
 221 import time
 222 class MyTimer:
 223         def __init__(self):
 224                 self.start = time.time()
 225
 226         def end(self):
 227                 self.end = time.time()
 228                 t = self.end-self.start
 229                 return t
 230
 231         def diff(self):
 232                 self.end = time.time()
 233                 t = self.end-self.start
 234                 self.start = self.end
 235                 return t