Minor code cleanup:
[OpenSRF.git] / 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 from osrf.system import osrfConnect
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 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     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 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.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" + osrf.json.osrfDebugNetworkObject(resp.content())
133         else:
134             print osrf.json.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     try:
225         plugins = osrfConfigValue('plugins')
226
227     except:
228         # XXX standard srfsh.xml does not yet define <plugins> element
229         print_red("No plugins defined in /srfsh/plugins/plugin\n")
230         return
231
232     plugins = osrfConfigValue('plugins.plugin')
233     if not isinstance(plugins, list):
234         plugins = [plugins]
235
236     for module in plugins:
237         name = module['module']
238         init = module['init']
239         print_green("Loading module %s..." % name)
240
241         try:
242             string = 'from %s import %s\n%s()' % (name, init, init)
243             exec(string)
244             print_red('OK\n')
245
246         except Exception, e:
247             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
248
249 def set_vars():
250     if not get_var('SRFSH_OUTPUT'):
251         set_var('SRFSH_OUTPUT', 'pretty')
252
253
254 def set_var(key, val):
255     os.environ[key] = val
256
257
258 def get_var(key):
259     try: return os.environ[key]
260     except: return ''
261     
262     
263 def print_green(string):
264     sys.stdout.write("\033[01;32m")
265     sys.stdout.write(string)
266     sys.stdout.write("\033[00m")
267     sys.stdout.flush()
268
269 def print_red(string):
270     sys.stdout.write("\033[01;31m")
271     sys.stdout.write(string)
272     sys.stdout.write("\033[00m")
273     sys.stdout.flush()
274
275
276
277
278 # Kick it off
279 set_vars()
280 setup_readline()
281 do_connect()
282 load_plugins()
283 do_loop()
284
285
286