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