Another potentially annoying typo
[OpenSRF.git] / src / python / srfsh.py
1 #!/usr/bin/python
2 # vim:et:ts=4
3 """
4 srfsh.py - provides a basic shell for issuing OpenSRF requests
5
6   help
7     - show this menu
8
9   math_bench <count>
10     - runs <count> opensrf.math requests and reports the average time
11
12   request <service> <method> [<param1>, <param2>, ...]
13     - performs an opensrf request
14
15   set VAR=<value>
16     - sets an environment variable
17
18   Environment variables:
19     SRFSH_OUTPUT = pretty - print pretty JSON and key/value pairs for network objects
20                  = raw - print formatted JSON 
21
22     SRFSH_LOCALE = <locale> - request responses to be returned in locale <locale> if available
23 """
24
25 import os, sys, time, readline, atexit, re
26 import osrf.json
27 import osrf.system
28 import osrf.ses
29 import osrf.conf
30 import osrf.log
31
32 # -------------------------------------------------------------------
33 # main listen loop
34 # -------------------------------------------------------------------
35 def do_loop():
36     while True:
37
38         try:
39             #line = raw_input("srfsh% ")
40             line = raw_input("\033[01;32msrfsh\033[01;34m% \033[00m")
41             if not len(line): 
42                 continue
43             if str.lower(line) == 'exit' or str.lower(line) == 'quit': 
44                 break
45             parts = str.split(line)
46
47             command = parts[0]
48         
49             if command == 'request':
50                 parts.pop(0)
51                 handle_request(parts)
52                 continue
53
54             if command == 'math_bench':
55                 parts.pop(0)
56                 handle_math_bench(parts)
57                 continue
58
59             if command == 'help':
60                 handle_help()
61                 continue
62
63             if command == 'set':
64                 parts.pop(0)
65                 handle_set(parts)
66
67             if command == 'get':
68                 parts.pop(0)
69                 handle_get(parts)
70
71
72
73         except KeyboardInterrupt:
74             print ""
75
76         except EOFError:
77             print "exiting..."
78             sys.exit(0)
79
80
81 # -------------------------------------------------------------------
82 # Set env variables to control behavior
83 # -------------------------------------------------------------------
84 def handle_set(parts):
85     cmd = "".join(parts)
86     pattern = re.compile('(.*)=(.*)').match(cmd)
87     key = pattern.group(1)
88     val = pattern.group(2)
89     set_var(key, val)
90     print "%s = %s" % (key, val)
91
92 def handle_get(parts):
93     try:
94         print get_var(parts[0])
95     except:
96         print ""
97
98
99 # -------------------------------------------------------------------
100 # Prints help info
101 # -------------------------------------------------------------------
102 def handle_help():
103     print __doc__
104
105 # -------------------------------------------------------------------
106 # performs an opensrf request
107 # -------------------------------------------------------------------
108 def handle_request(parts):
109     service = parts.pop(0)
110     method = parts.pop(0)
111     locale = __get_locale()
112     jstr = '[%s]' % "".join(parts)
113     params = None
114
115     try:
116         params = osrf.json.to_object(jstr)
117     except:
118         print "Error parsing JSON: %s" % jstr
119         return
120
121     ses = osrf.ses.ClientSession(service, locale=locale)
122
123     start = time.time()
124
125     req = ses.request2(method, tuple(params))
126
127
128     while True:
129         resp = req.recv(timeout=120)
130         osrf.log.log_internal("Looping through receive request")
131         if not resp:
132             break
133         total = time.time() - start
134
135         otp = get_var('SRFSH_OUTPUT')
136         if otp == 'pretty':
137             print "\n" + osrf.json.debug_net_object(resp.content())
138         else:
139             print osrf.json.pprint(osrf.json.to_json(resp.content()))
140
141     req.cleanup()
142     ses.cleanup()
143
144     print '-'*60
145     print "Total request time: %f" % total
146     print '-'*60
147
148
149 def handle_math_bench(parts):
150
151     count = int(parts.pop(0))
152     ses = osrf.ses.ClientSession('opensrf.math')
153     times = []
154
155     for cnt in range(100):
156         if cnt % 10:
157             sys.stdout.write('.')
158         else:
159             sys.stdout.write( str( cnt / 10 ) )
160     print ""
161
162
163     for cnt in range(count):
164     
165         starttime = time.time()
166         req = ses.request('add', 1, 2)
167         resp = req.recv(timeout=2)
168         endtime = time.time()
169     
170         if resp.content() == 3:
171             sys.stdout.write("+")
172             sys.stdout.flush()
173             times.append( endtime - starttime )
174         else:
175             print "What happened? %s" % str(resp.content())
176     
177         req.cleanup()
178         if not ( (cnt + 1) % 100):
179             print ' [%d]' % (cnt + 1)
180     
181     ses.cleanup()
182     total = 0
183     for cnt in times:
184         total += cnt 
185     print "\naverage time %f" % (total / len(times))
186
187
188
189
190 # -------------------------------------------------------------------
191 # Defines the tab-completion handling and sets up the readline history 
192 # -------------------------------------------------------------------
193 def setup_readline():
194     class SrfshCompleter(object):
195         def __init__(self, words):
196             self.words = words
197             self.prefix = None
198     
199         def complete(self, prefix, index):
200             if prefix != self.prefix:
201                 # find all words that start with this prefix
202                 self.matching_words = [
203                     w for w in self.words if w.startswith(prefix)
204                 ]
205                 self.prefix = prefix
206                 try:
207                     return self.matching_words[index]
208                 except IndexError:
209                     return None
210     
211     words = 'request', 'help', 'exit', 'quit', 'opensrf.settings', 'opensrf.math', 'set'
212     completer = SrfshCompleter(words)
213     readline.parse_and_bind("tab: complete")
214     readline.set_completer(completer.complete)
215
216     histfile = os.path.join(get_var('HOME'), ".srfsh_history")
217     try:
218         readline.read_history_file(histfile)
219     except IOError:
220         pass
221     atexit.register(readline.write_history_file, histfile)
222
223 def do_connect():
224     file = os.path.join(get_var('HOME'), ".srfsh.xml")
225     print_green("Connecting to opensrf...")
226     osrf.system.System.connect(config_file=file, config_context='srfsh')
227     print_red('OK\n')
228
229 def load_plugins():
230     # Load the user defined external plugins
231     # XXX Make this a real module interface, with tab-complete words, commands, etc.
232     try:
233         plugins = osrf.conf.get('plugins')
234
235     except:
236         # XXX standard srfsh.xml does not yet define <plugins> element
237         print_red("No plugins defined in /srfsh/plugins/plugin\n")
238         return
239
240     plugins = osrf.conf.get('plugins.plugin')
241     if not isinstance(plugins, list):
242         plugins = [plugins]
243
244     for module in plugins:
245         name = module['module']
246         init = module['init']
247         print_green("Loading module %s..." % name)
248
249         try:
250             string = 'import %s\n%s.%s()' % (name, name, init)
251             exec(string)
252             print_red('OK\n')
253
254         except Exception, e:
255             sys.stderr.write("\nError importing plugin %s, with init symbol %s: \n%s\n" % (name, init, e))
256
257 def set_vars():
258     if not get_var('SRFSH_OUTPUT'):
259         set_var('SRFSH_OUTPUT', 'pretty')
260
261     # XXX Do we need to differ between LANG and LC_MESSAGES?
262     if not get_var('SRFSH_LOCALE'):
263         set_var('SRFSH_LOCALE', get_var('LC_ALL'))
264
265 def set_var(key, val):
266     os.environ[key] = val
267
268 def get_var(key):
269     return os.environ.get(key, '')
270     
271 def __get_locale():
272     """
273     Return the defined locale for this srfsh session.
274
275     A locale in OpenSRF is currently defined as a [a-z]{2}-[A-Z]{2} pattern.
276     This function munges the LC_ALL setting to conform to that pattern; for
277     example, trimming en_CA.UTF-8 to en-CA.
278
279     >>> import srfsh
280     >>> srfsh.set_var('SRFSH_LOCALE', 'zz-ZZ')
281     >>> print __get_locale()
282     zz-ZZ
283     >>> srfsh.set_var('SRFSH_LOCALE', 'en_CA.UTF-8')
284     >>> print __get_locale()
285     en-CA
286     """
287
288     env_locale = get_var('SRFSH_LOCALE')
289     if env_locale:
290         pattern = re.compile(r'^\s*([a-z]+)[^a-zA-Z]([A-Z]+)').search(env_locale)
291         lang = pattern.group(1)
292         region = pattern.group(2)
293         locale = "%s-%s" % (lang, region)
294     else:
295         locale = 'en-US'
296
297     return locale
298     
299 def print_green(string):
300     sys.stdout.write("\033[01;32m")
301     sys.stdout.write(string)
302     sys.stdout.write("\033[00m")
303     sys.stdout.flush()
304
305 def print_red(string):
306     sys.stdout.write("\033[01;31m")
307     sys.stdout.write(string)
308     sys.stdout.write("\033[00m")
309     sys.stdout.flush()
310
311
312 if __name__ == '__main__':
313
314     # Kick it off
315     set_vars()
316     setup_readline()
317     do_connect()
318     load_plugins()
319     do_loop()
320