src/python/srfsh.py

   1 #!/usr/bin/python2.4
   2 import os, sys, time, readline, atexit, re
   3 from string import *
   4 from osrf.system import osrfConnect
   5 from osrf.json import *
   6 from osrf.ses import osrfClientSession
   7 from osrf.conf import osrfConfigValue
   8
   9
  10 # -------------------------------------------------------------------
  11 # main listen loop
  12 # -------------------------------------------------------------------
  13 def do_loop():
  14         while True:
  15
  16                 try:
  17                         #line = raw_input("srfsh% ")
  18                         line = raw_input("\033[01;32msrfsh\033[01;34m% \033[00m")
  19                         if not len(line):
  20                                 continue
  21                         if lower(line) == 'exit' or lower(line) == 'quit':
  22                                 break
  23                         parts = split(line)
  24
  25                         command = parts[0]
  26
  27                         if command == 'request':
  28                                 parts.pop(0)
  29                                 handle_request(parts)
  30                                 continue
  31
  32                         if command == 'math_bench':
  33                                 parts.pop(0)
  34                                 handle_math_bench(parts)
  35                                 continue
  36
  37                         if command == 'help':
  38                                 handle_help()
  39                                 continue
  40
  41                         if command == 'set':
  42                                 parts.pop(0)
  43                                 handle_set(parts)
  44
  45                         if command == 'get':
  46                                 parts.pop(0)
  47                                 handle_get(parts)
  48
  49
  50
  51                 except KeyboardInterrupt:
  52                         print ""
  53
  54                 except EOFError:
  55                         print "exiting..."
  56                         sys.exit(0)
  57
  58
  59 # -------------------------------------------------------------------
  60 # Set env variables to control behavior
  61 # -------------------------------------------------------------------
  62 def handle_set(parts):
  63         m = re.compile('(.*)=(.*)').match(parts[0])
  64         key = m.group(1)
  65         val = m.group(2)
  66         set_var(key, val)
  67         print "%s = %s" % (key, val)
  68
  69 def handle_get(parts):
  70         try:
  71                 print get_var(parts[0])
  72         except:
  73                 print ""
  74
  75
  76 # -------------------------------------------------------------------
  77 # Prints help info
  78 # -------------------------------------------------------------------
  79 def handle_help():
  80         print """
  81   help
  82     - show this menu
  83
  84   math_bench <count>
  85     - runs <count> opensrf.math requests and reports the average time
  86
  87   request <service> <method> [<param1>, <param2>, ...]
  88     - performs an opensrf request
  89
  90   set VAR=<value>
  91     - sets an environment variable
  92
  93   Environment variables:
  94     SRFSH_OUTPUT = pretty - print pretty JSON and key/value pairs for network objects
  95                  = raw - print formatted JSON
  96         """
  97
  98
  99
 100
 101 # -------------------------------------------------------------------
 102 # performs an opesnrf request
 103 # -------------------------------------------------------------------
 104 def handle_request(parts):
 105         service = parts.pop(0)
 106         method = parts.pop(0)
 107         jstr = '[%s]' % join(parts)
 108         params = None
 109
 110         try:
 111                 params = osrfJSONToObject(jstr)
 112         except:
 113                 print "Error parsing JSON: %s" % jstr
 114                 return
 115
 116         ses = osrfClientSession(service)
 117
 118         end = None
 119         start = time.time()
 120
 121         req = ses.request2(method, tuple(params))
 122
 123
 124         while True:
 125                 resp = req.recv(timeout=120)
 126                 if not end:
 127                         total = time.time() - start
 128                 if not resp: break
 129
 130                 otp = get_var('SRFSH_OUTPUT')
 131                 if otp == 'pretty':
 132                         print "\n" + osrfDebugNetworkObject(resp.content())
 133                 else:
 134                         print osrfFormatJSON(osrfObjectToJSON(resp.content()))
 135
 136         req.cleanup()
 137         ses.cleanup()
 138
 139         print '-'*60
 140         print "Total request time: %f" % total
 141         print '-'*60
 142
 143
 144 def handle_math_bench(parts):
 145
 146         count = int(parts.pop(0))
 147         ses = osrfClientSession('opensrf.math')
 148         times = []
 149
 150         for i in range(100):
 151                 if i % 10: sys.stdout.write('.')
 152                 else: sys.stdout.write( str( i / 10 ) )
 153         print "";
 154
 155
 156         for i in range(count):
 157
 158                 starttime = time.time()
 159                 req = ses.request('add', 1, 2)
 160                 resp = req.recv(timeout=2)
 161                 endtime = time.time()
 162
 163                 if resp.content() == 3:
 164                         sys.stdout.write("+")
 165                         sys.stdout.flush()
 166                         times.append( endtime - starttime )
 167                 else:
 168                         print "What happened? %s" % str(resp.content())
 169
 170                 req.cleanup()
 171                 if not ( (i+1) % 100):
 172                         print ' [%d]' % (i+1)
 173
 174         ses.cleanup()
 175         total = 0
 176         for i in times: total += i
 177         print "\naverage time %f" % (total / len(times))
 178
 179
 180
 181
 182 # -------------------------------------------------------------------
 183 # Defines the tab-completion handling and sets up the readline history
 184 # -------------------------------------------------------------------
 185 def setup_readline():
 186         class SrfshCompleter(object):
 187                 def __init__(self, words):
 188                         self.words = words
 189                         self.prefix = None
 190
 191                 def complete(self, prefix, index):
 192                         if prefix != self.prefix:
 193                                 # find all words that start with this prefix
 194                                 self.matching_words = [
 195                                         w for w in self.words if w.startswith(prefix)
 196                                 ]
 197                                 self.prefix = prefix
 198                                 try:
 199                                         return self.matching_words[index]
 200                                 except IndexError:
 201                                         return None
 202
 203         words = 'request', 'help', 'exit', 'quit', 'opensrf.settings', 'opensrf.math', 'set'
 204         completer = SrfshCompleter(words)
 205         readline.parse_and_bind("tab: complete")
 206         readline.set_completer(completer.complete)
 207
 208         histfile = os.path.join(get_var('HOME'), ".srfsh_history")
 209         try:
 210             readline.read_history_file(histfile)
 211         except IOError:
 212                 pass
 213         atexit.register(readline.write_history_file, histfile)
 214
 215 def do_connect():
 216         file = os.path.join(get_var('HOME'), ".srfsh.xml")
 217         print_green("Connecting to opensrf...")
 218         osrfConnect(file, 'srfsh')
 219         print_red('OK\n')
 220
 221 def load_plugins():
 222         # Load the user defined external plugins
 223         # XXX Make this a real module interface, with tab-complete words, commands, etc.
 224         plugins = osrfConfigValue('plugins')
 225         plugins = osrfConfigValue('plugins.plugin')
 226         if not isinstance(plugins, list):
 227                 plugins = [plugins]
 228
 229         for module in plugins:
 230                 name = module['module']
 231                 init = module['init']
 232                 print_green("Loading module %s..." % name)
 233
 234                 try:
 235                         str = 'from %s import %s\n%s()' % (name, init, init)
 236                         exec(str)
 237                         print_red('OK\n')
 238
 239                 except Exception, e:
 240                         sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
 241
 242 def set_vars():
 243         if not get_var('SRFSH_OUTPUT'):
 244                 set_var('SRFSH_OUTPUT', 'pretty')
 245
 246
 247 def set_var(key, val):
 248         os.environ[key] = val
 249
 250
 251 def get_var(key):
 252         try: return os.environ[key]
 253         except: return ''
 254
 255
 256 def print_green(str):
 257         sys.stdout.write("\033[01;32m")
 258         sys.stdout.write(str)
 259         sys.stdout.write("\033[00m")
 260         sys.stdout.flush()
 261
 262 def print_red(str):
 263         sys.stdout.write("\033[01;31m")
 264         sys.stdout.write(str)
 265         sys.stdout.write("\033[00m")
 266         sys.stdout.flush()
 267
 268
 269
 270
 271 # Kick it off
 272 set_vars()
 273 setup_readline()
 274 do_connect()
 275 load_plugins()
 276 do_loop()
 277
 278
 279