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 import osrf.log
   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 str.lower(line) == 'exit' or str.lower(line) == 'quit':
  23                 break
  24             parts = str.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     pattern = re.compile('(.*)=(.*)').match(parts[0])
  65     key = pattern.group(1)
  66     val = pattern.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 = osrf.json.to_object(jstr)
 113     except:
 114         print "Error parsing JSON: %s" % jstr
 115         return
 116
 117     ses = osrf.ses.ClientSession(service)
 118
 119     start = time.time()
 120
 121     req = ses.request2(method, tuple(params))
 122
 123
 124     while True:
 125         resp = req.recv(timeout=120)
 126         osrf.log.log_internal("Looping through receive request")
 127         if not resp:
 128             break
 129         total = time.time() - start
 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