Add some basic locale support.
[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 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