src/python/srfsh.py

   1 #!/usr/bin/python
   2 # vim:et:ts=4
   3 import os, sys, time, readline, atexit, re
   4 import osrf.json
   5 import osrf.system
   6 import osrf.ses
   7 import osrf.conf
   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 str.lower(line) == 'exit' or str.lower(line) == 'quit':
  22                 break
  23             parts = str.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     pattern = re.compile('(.*)=(.*)').match(parts[0])
  64     key = pattern.group(1)
  65     val = pattern.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 opensrf 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 = osrf.json.to_object(jstr)
 112     except:
 113         print "Error parsing JSON: %s" % jstr
 114         return
 115
 116     ses = osrf.ses.ClientSession(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:
 129             break
 130
 131         otp = get_var('SRFSH_OUTPUT')
 132         if otp == 'pretty':
 133             print "\n" + osrf.json.debug_net_object(resp.content())
 134         else:
 135             print osrf.json.pprint(osrf.json.to_json(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 = osrf.ses.ClientSession('opensrf.math')
 149     times = []
 150
 151     for cnt in range(100):
 152         if cnt % 10:
 153             sys.stdout.write('.')
 154         else:
 155             sys.stdout.write( str( cnt / 10 ) )
 156     print ""
 157
 158
 159     for cnt in range(count):
 160
 161         starttime = time.time()
 162         req = ses.request('add', 1, 2)
 163         resp = req.recv(timeout=2)
 164         endtime = time.time()
 165
 166         if resp.content() == 3:
 167             sys.stdout.write("+")
 168             sys.stdout.flush()
 169             times.append( endtime - starttime )
 170         else:
 171             print "What happened? %s" % str(resp.content())
 172
 173         req.cleanup()
 174         if not ( (cnt + 1) % 100):
 175             print ' [%d]' % (cnt + 1)
 176
 177     ses.cleanup()
 178     total = 0
 179     for cnt in times:
 180         total += cnt
 181     print "\naverage time %f" % (total / len(times))
 182
 183
 184
 185
 186 # -------------------------------------------------------------------
 187 # Defines the tab-completion handling and sets up the readline history
 188 # -------------------------------------------------------------------
 189 def setup_readline():
 190     class SrfshCompleter(object):
 191         def __init__(self, words):
 192             self.words = words
 193             self.prefix = None
 194
 195         def complete(self, prefix, index):
 196             if prefix != self.prefix:
 197                 # find all words that start with this prefix
 198                 self.matching_words = [
 199                     w for w in self.words if w.startswith(prefix)
 200                 ]
 201                 self.prefix = prefix
 202                 try:
 203                     return self.matching_words[index]
 204                 except IndexError:
 205                     return None
 206
 207     words = 'request', 'help', 'exit', 'quit', 'opensrf.settings', 'opensrf.math', 'set'
 208     completer = SrfshCompleter(words)
 209     readline.parse_and_bind("tab: complete")
 210     readline.set_completer(completer.complete)
 211
 212     histfile = os.path.join(get_var('HOME'), ".srfsh_history")
 213     try:
 214         readline.read_history_file(histfile)
 215     except IOError:
 216         pass
 217     atexit.register(readline.write_history_file, histfile)
 218
 219 def do_connect():
 220     file = os.path.join(get_var('HOME'), ".srfsh.xml")
 221     print_green("Connecting to opensrf...")
 222     osrf.system.connect(file, 'srfsh')
 223     print_red('OK\n')
 224
 225 def load_plugins():
 226     # Load the user defined external plugins
 227     # XXX Make this a real module interface, with tab-complete words, commands, etc.
 228     try:
 229         plugins = osrf.conf.get('plugins')
 230
 231     except:
 232         # XXX standard srfsh.xml does not yet define <plugins> element
 233         print_red("No plugins defined in /srfsh/plugins/plugin\n")
 234         return
 235
 236     plugins = osrf.conf.get('plugins.plugin')
 237     if not isinstance(plugins, list):
 238         plugins = [plugins]
 239
 240     for module in plugins:
 241         name = module['module']
 242         init = module['init']
 243         print_green("Loading module %s..." % name)
 244
 245         try:
 246             string = 'from %s import %s\n%s()' % (name, init, init)
 247             exec(string)
 248             print_red('OK\n')
 249
 250         except Exception, e:
 251             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
 252
 253 def set_vars():
 254     if not get_var('SRFSH_OUTPUT'):
 255         set_var('SRFSH_OUTPUT', 'pretty')
 256
 257
 258 def set_var(key, val):
 259     os.environ[key] = val
 260
 261
 262 def get_var(key):
 263     try:
 264         return os.environ[key]
 265     except:
 266         return ''
 267
 268
 269 def print_green(string):
 270     sys.stdout.write("\033[01;32m")
 271     sys.stdout.write(string)
 272     sys.stdout.write("\033[00m")
 273     sys.stdout.flush()
 274
 275 def print_red(string):
 276     sys.stdout.write("\033[01;31m")
 277     sys.stdout.write(string)
 278     sys.stdout.write("\033[00m")
 279     sys.stdout.flush()
 280
 281
 282
 283
 284 # Kick it off
 285 set_vars()
 286 setup_readline()
 287 do_connect()
 288 load_plugins()
 289 do_loop()
 290
 291
 292