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(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(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(name):
  68                                 debugname = "debug.%s" % name
  69                                 if not self.exists(debugname):
  70                                         shutil.copyfile(self.__file(name), self.__file(debugname))
  71                                 name = debugname
  72                         else:   # neither exist
  73                                 raise Exception, "No such pickle based on %s" % self.__file(name)
  74                 else:
  75                         if not self.exists(name):
  76                                 raise Exception, "No such file %s" % name
  77
  78                 print "loading %s" % self.__file(name)
  79                 f = open(self.__file(name), 'r')
  80                 o = pickle.load(f)
  81                 f.close()
  82                 return o
  83
  84
  85         # use the environment to extract the data associated with the local
  86         # variable 'name'
  87         def dump(self, name, obj=None, depth=1):
  88                 if obj == None:
  89                         o = inspect.getouterframes(inspect.currentframe())
  90                         up1 = o[depth][0] # get the frame one prior to (up from) this frame
  91                         argvals = inspect.getargvalues(up1)
  92                         # TODO: check that 'name' is a local variable; otherwise this would fail.
  93                         obj = argvals[3][name] # extract the local variable name 'name'
  94                 if not os.path.isdir("%s/" % PICKLE_PATH):
  95                         os.mkdir("%s" % PICKLE_PATH)
  96                 if self.config.debug:
  97                         name = "debug.%s" % name
  98                 f = open(self.__file(name), 'w')
  99                 pickle.dump(obj, f)
 100                 f.close()
 101                 return
 102
 103
 104 ssh_options = { 'StrictHostKeyChecking':'no',
 105                                 'BatchMode':'yes',
 106                                 'PasswordAuthentication':'no',
 107                                 'ConnectTimeout':'20'}
 108
 109 class SSH:
 110         def __init__(self, user, host, options = ssh_options):
 111                 self.options = options
 112                 self.user = user
 113                 self.host = host
 114                 return
 115
 116         def __options_to_str(self):
 117                 options = ""
 118                 for o,v in self.options.iteritems():
 119                         options = options + "-o %s=%s " % (o,v)
 120                 return options
 121
 122         def run(self, cmd):
 123                 cmd = "ssh %s %s@%s '%s'" % (self.__options_to_str(),
 124                                                                         self.user, self.host, cmd)
 125                 if ( DEBUG == 1 ):
 126                         print cmd,
 127                 (f_in, f_out, f_err) = os.popen3(cmd)
 128                 value = f_out.read()
 129                 if value == "":
 130                         raise Exception, f_err.read()
 131                 if ( DEBUG == 1 ):
 132                         print " == %s" % value
 133                 f_out.close()
 134                 f_in.close()
 135                 f_err.close()
 136                 return value
 137
 138         def runE(self, cmd):
 139                 cmd = "ssh %s %s@%s '%s'" % (self.__options_to_str(),
 140                                                                         self.user, self.host, cmd)
 141                 if ( DEBUG == 1 ):
 142                         print cmd,
 143                 (f_in, f_out, f_err) = os.popen3(cmd)
 144
 145                 value = f_out.read()
 146                 if value == "": # An error has occured
 147                         value = f_err.read()
 148
 149                 if ( DEBUG == 1 ):
 150                         print " == %s" % value
 151                 f_out.close()
 152                 f_in.close()
 153                 f_err.close()
 154                 return value.strip()
 155
 156 import time
 157 class MyTimer:
 158         def __init__(self):
 159                 self.start = time.time()
 160
 161         def end(self):
 162                 self.end = time.time()
 163                 t = self.end-self.start
 164                 return t
 165
 166         def diff(self):
 167                 self.end = time.time()
 168                 t = self.end-self.start
 169                 self.start = self.end
 170                 return t