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