bb058b143349c5bb4f798890e2052c4e2f7654ba
[OpenSRF.git] / 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