src/python/srfsh.py

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