~ipython-contrib/+junk/ipython-zmq

« back to all changes in this revision

Viewing changes to IPython/irunner.py

  • Committer: ville
  • Date: 2008-02-16 09:50:47 UTC
  • mto: (0.12.1 ipython_main)
  • mto: This revision was merged to the branch mainline in revision 990.
  • Revision ID: ville@ville-pc-20080216095047-500x6dluki1iz40o
initialization (no svn history)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
#!/usr/bin/env python
 
2
"""Module for interactively running scripts.
 
3
 
 
4
This module implements classes for interactively running scripts written for
 
5
any system with a prompt which can be matched by a regexp suitable for
 
6
pexpect.  It can be used to run as if they had been typed up interactively, an
 
7
arbitrary series of commands for the target system.
 
8
 
 
9
The module includes classes ready for IPython (with the default prompts),
 
10
plain Python and SAGE, but making a new one is trivial.  To see how to use it,
 
11
simply run the module as a script:
 
12
 
 
13
./irunner.py --help
 
14
 
 
15
 
 
16
This is an extension of Ken Schutte <kschutte-AT-csail.mit.edu>'s script
 
17
contributed on the ipython-user list:
 
18
 
 
19
http://scipy.net/pipermail/ipython-user/2006-May/001705.html
 
20
 
 
21
 
 
22
NOTES:
 
23
 
 
24
 - This module requires pexpect, available in most linux distros, or which can
 
25
 be downloaded from
 
26
 
 
27
    http://pexpect.sourceforge.net
 
28
 
 
29
 - Because pexpect only works under Unix or Windows-Cygwin, this has the same
 
30
 limitations.  This means that it will NOT work under native windows Python.
 
31
"""
 
32
 
 
33
# Stdlib imports
 
34
import optparse
 
35
import os
 
36
import sys
 
37
 
 
38
# Third-party modules.
 
39
import pexpect
 
40
 
 
41
# Global usage strings, to avoid indentation issues when typing it below.
 
42
USAGE = """
 
43
Interactive script runner, type: %s
 
44
 
 
45
runner [opts] script_name
 
46
"""
 
47
 
 
48
def pexpect_monkeypatch():
 
49
    """Patch pexpect to prevent unhandled exceptions at VM teardown.
 
50
 
 
51
    Calling this function will monkeypatch the pexpect.spawn class and modify
 
52
    its __del__ method to make it more robust in the face of failures that can
 
53
    occur if it is called when the Python VM is shutting down.
 
54
 
 
55
    Since Python may fire __del__ methods arbitrarily late, it's possible for
 
56
    them to execute during the teardown of the Python VM itself.  At this
 
57
    point, various builtin modules have been reset to None.  Thus, the call to
 
58
    self.close() will trigger an exception because it tries to call os.close(),
 
59
    and os is now None.
 
60
    """
 
61
    
 
62
    if pexpect.__version__[:3] >= '2.2':
 
63
        # No need to patch, fix is already the upstream version.
 
64
        return
 
65
    
 
66
    def __del__(self):
 
67
        """This makes sure that no system resources are left open.
 
68
        Python only garbage collects Python objects. OS file descriptors
 
69
        are not Python objects, so they must be handled explicitly.
 
70
        If the child file descriptor was opened outside of this class
 
71
        (passed to the constructor) then this does not close it.
 
72
        """
 
73
        if not self.closed:
 
74
            try:
 
75
                self.close()
 
76
            except AttributeError:
 
77
                pass
 
78
 
 
79
    pexpect.spawn.__del__ = __del__
 
80
 
 
81
pexpect_monkeypatch()
 
82
 
 
83
# The generic runner class
 
84
class InteractiveRunner(object):
 
85
    """Class to run a sequence of commands through an interactive program."""
 
86
    
 
87
    def __init__(self,program,prompts,args=None,out=sys.stdout,echo=True):
 
88
        """Construct a runner.
 
89
 
 
90
        Inputs:
 
91
 
 
92
          - program: command to execute the given program.
 
93
 
 
94
          - prompts: a list of patterns to match as valid prompts, in the
 
95
          format used by pexpect.  This basically means that it can be either
 
96
          a string (to be compiled as a regular expression) or a list of such
 
97
          (it must be a true list, as pexpect does type checks).
 
98
 
 
99
          If more than one prompt is given, the first is treated as the main
 
100
          program prompt and the others as 'continuation' prompts, like
 
101
          python's.  This means that blank lines in the input source are
 
102
          ommitted when the first prompt is matched, but are NOT ommitted when
 
103
          the continuation one matches, since this is how python signals the
 
104
          end of multiline input interactively.
 
105
 
 
106
        Optional inputs:
 
107
 
 
108
          - args(None): optional list of strings to pass as arguments to the
 
109
          child program.
 
110
 
 
111
          - out(sys.stdout): if given, an output stream to be used when writing
 
112
          output.  The only requirement is that it must have a .write() method.
 
113
 
 
114
        Public members not parameterized in the constructor:
 
115
 
 
116
          - delaybeforesend(0): Newer versions of pexpect have a delay before
 
117
          sending each new input.  For our purposes here, it's typically best
 
118
          to just set this to zero, but if you encounter reliability problems
 
119
          or want an interactive run to pause briefly at each prompt, just
 
120
          increase this value (it is measured in seconds).  Note that this
 
121
          variable is not honored at all by older versions of pexpect.
 
122
        """
 
123
        
 
124
        self.program = program
 
125
        self.prompts = prompts
 
126
        if args is None: args = []
 
127
        self.args = args
 
128
        self.out = out
 
129
        self.echo = echo
 
130
        # Other public members which we don't make as parameters, but which
 
131
        # users may occasionally want to tweak
 
132
        self.delaybeforesend = 0
 
133
 
 
134
        # Create child process and hold on to it so we don't have to re-create
 
135
        # for every single execution call
 
136
        c = self.child = pexpect.spawn(self.program,self.args,timeout=None)
 
137
        c.delaybeforesend = self.delaybeforesend
 
138
        # pexpect hard-codes the terminal size as (24,80) (rows,columns).
 
139
        # This causes problems because any line longer than 80 characters gets
 
140
        # completely overwrapped on the printed outptut (even though
 
141
        # internally the code runs fine).  We reset this to 99 rows X 200
 
142
        # columns (arbitrarily chosen), which should avoid problems in all
 
143
        # reasonable cases.
 
144
        c.setwinsize(99,200)
 
145
 
 
146
    def close(self):
 
147
        """close child process"""
 
148
 
 
149
        self.child.close()
 
150
 
 
151
    def run_file(self,fname,interact=False,get_output=False):
 
152
        """Run the given file interactively.
 
153
 
 
154
        Inputs:
 
155
 
 
156
          -fname: name of the file to execute.
 
157
 
 
158
        See the run_source docstring for the meaning of the optional
 
159
        arguments."""
 
160
 
 
161
        fobj = open(fname,'r')
 
162
        try:
 
163
            out = self.run_source(fobj,interact,get_output)
 
164
        finally:
 
165
            fobj.close()
 
166
        if get_output:
 
167
            return out
 
168
 
 
169
    def run_source(self,source,interact=False,get_output=False):
 
170
        """Run the given source code interactively.
 
171
 
 
172
        Inputs:
 
173
 
 
174
          - source: a string of code to be executed, or an open file object we
 
175
          can iterate over.
 
176
 
 
177
        Optional inputs:
 
178
 
 
179
          - interact(False): if true, start to interact with the running
 
180
          program at the end of the script.  Otherwise, just exit.
 
181
 
 
182
          - get_output(False): if true, capture the output of the child process
 
183
          (filtering the input commands out) and return it as a string.
 
184
 
 
185
        Returns:
 
186
          A string containing the process output, but only if requested.
 
187
          """
 
188
 
 
189
        # if the source is a string, chop it up in lines so we can iterate
 
190
        # over it just as if it were an open file.
 
191
        if not isinstance(source,file):
 
192
            source = source.splitlines(True)
 
193
 
 
194
        if self.echo:
 
195
            # normalize all strings we write to use the native OS line
 
196
            # separators.
 
197
            linesep  = os.linesep
 
198
            stdwrite = self.out.write
 
199
            write    = lambda s: stdwrite(s.replace('\r\n',linesep))
 
200
        else:
 
201
            # Quiet mode, all writes are no-ops
 
202
            write = lambda s: None
 
203
            
 
204
        c = self.child
 
205
        prompts = c.compile_pattern_list(self.prompts)
 
206
        prompt_idx = c.expect_list(prompts)
 
207
 
 
208
        # Flag whether the script ends normally or not, to know whether we can
 
209
        # do anything further with the underlying process.
 
210
        end_normal = True
 
211
 
 
212
        # If the output was requested, store it in a list for return at the end
 
213
        if get_output:
 
214
            output = []
 
215
            store_output = output.append
 
216
        
 
217
        for cmd in source:
 
218
            # skip blank lines for all matches to the 'main' prompt, while the
 
219
            # secondary prompts do not
 
220
            if prompt_idx==0 and \
 
221
                   (cmd.isspace() or cmd.lstrip().startswith('#')):
 
222
                write(cmd)
 
223
                continue
 
224
 
 
225
            #write('AFTER: '+c.after)  # dbg
 
226
            write(c.after)
 
227
            c.send(cmd)
 
228
            try:
 
229
                prompt_idx = c.expect_list(prompts)
 
230
            except pexpect.EOF:
 
231
                # this will happen if the child dies unexpectedly
 
232
                write(c.before)
 
233
                end_normal = False
 
234
                break
 
235
            
 
236
            write(c.before)
 
237
 
 
238
            # With an echoing process, the output we get in c.before contains
 
239
            # the command sent, a newline, and then the actual process output
 
240
            if get_output:
 
241
                store_output(c.before[len(cmd+'\n'):])
 
242
                #write('CMD: <<%s>>' % cmd)  # dbg
 
243
                #write('OUTPUT: <<%s>>' % output[-1])  # dbg
 
244
 
 
245
        self.out.flush()
 
246
        if end_normal:
 
247
            if interact:
 
248
                c.send('\n')
 
249
                print '<< Starting interactive mode >>',
 
250
                try:
 
251
                    c.interact()
 
252
                except OSError:
 
253
                    # This is what fires when the child stops.  Simply print a
 
254
                    # newline so the system prompt is aligned.  The extra
 
255
                    # space is there to make sure it gets printed, otherwise
 
256
                    # OS buffering sometimes just suppresses it.
 
257
                    write(' \n')
 
258
                    self.out.flush()
 
259
        else:
 
260
            if interact:
 
261
                e="Further interaction is not possible: child process is dead."
 
262
                print >> sys.stderr, e
 
263
 
 
264
        # Leave the child ready for more input later on, otherwise select just
 
265
        # hangs on the second invocation.
 
266
        c.send('\n')
 
267
        
 
268
        # Return any requested output
 
269
        if get_output:
 
270
            return ''.join(output)
 
271
                
 
272
    def main(self,argv=None):
 
273
        """Run as a command-line script."""
 
274
 
 
275
        parser = optparse.OptionParser(usage=USAGE % self.__class__.__name__)
 
276
        newopt = parser.add_option
 
277
        newopt('-i','--interact',action='store_true',default=False,
 
278
               help='Interact with the program after the script is run.')
 
279
 
 
280
        opts,args = parser.parse_args(argv)
 
281
 
 
282
        if len(args) != 1:
 
283
            print >> sys.stderr,"You must supply exactly one file to run."
 
284
            sys.exit(1)
 
285
 
 
286
        self.run_file(args[0],opts.interact)
 
287
 
 
288
 
 
289
# Specific runners for particular programs
 
290
class IPythonRunner(InteractiveRunner):
 
291
    """Interactive IPython runner.
 
292
 
 
293
    This initalizes IPython in 'nocolor' mode for simplicity.  This lets us
 
294
    avoid having to write a regexp that matches ANSI sequences, though pexpect
 
295
    does support them.  If anyone contributes patches for ANSI color support,
 
296
    they will be welcome.
 
297
 
 
298
    It also sets the prompts manually, since the prompt regexps for
 
299
    pexpect need to be matched to the actual prompts, so user-customized
 
300
    prompts would break this.
 
301
    """
 
302
    
 
303
    def __init__(self,program = 'ipython',args=None,out=sys.stdout,echo=True):
 
304
        """New runner, optionally passing the ipython command to use."""
 
305
        
 
306
        args0 = ['-colors','NoColor',
 
307
                 '-pi1','In [\\#]: ',
 
308
                 '-pi2','   .\\D.: ',
 
309
                 '-noterm_title',
 
310
                 '-noautoindent']
 
311
        if args is None: args = args0
 
312
        else: args = args0 + args
 
313
        prompts = [r'In \[\d+\]: ',r'   \.*: ']
 
314
        InteractiveRunner.__init__(self,program,prompts,args,out,echo)
 
315
 
 
316
 
 
317
class PythonRunner(InteractiveRunner):
 
318
    """Interactive Python runner."""
 
319
 
 
320
    def __init__(self,program='python',args=None,out=sys.stdout,echo=True):
 
321
        """New runner, optionally passing the python command to use."""
 
322
 
 
323
        prompts = [r'>>> ',r'\.\.\. ']
 
324
        InteractiveRunner.__init__(self,program,prompts,args,out,echo)
 
325
 
 
326
 
 
327
class SAGERunner(InteractiveRunner):
 
328
    """Interactive SAGE runner.
 
329
    
 
330
    WARNING: this runner only works if you manually configure your SAGE copy
 
331
    to use 'colors NoColor' in the ipythonrc config file, since currently the
 
332
    prompt matching regexp does not identify color sequences."""
 
333
 
 
334
    def __init__(self,program='sage',args=None,out=sys.stdout,echo=True):
 
335
        """New runner, optionally passing the sage command to use."""
 
336
 
 
337
        prompts = ['sage: ',r'\s*\.\.\. ']
 
338
        InteractiveRunner.__init__(self,program,prompts,args,out,echo)
 
339
 
 
340
 
 
341
class RunnerFactory(object):
 
342
    """Code runner factory.
 
343
 
 
344
    This class provides an IPython code runner, but enforces that only one
 
345
    runner is ever instantiated.  The runner is created based on the extension
 
346
    of the first file to run, and it raises an exception if a runner is later
 
347
    requested for a different extension type.
 
348
 
 
349
    This ensures that we don't generate example files for doctest with a mix of
 
350
    python and ipython syntax.
 
351
    """
 
352
 
 
353
    def __init__(self,out=sys.stdout):
 
354
        """Instantiate a code runner."""
 
355
        
 
356
        self.out = out
 
357
        self.runner = None
 
358
        self.runnerClass = None
 
359
 
 
360
    def _makeRunner(self,runnerClass):
 
361
        self.runnerClass = runnerClass
 
362
        self.runner = runnerClass(out=self.out)
 
363
        return self.runner
 
364
          
 
365
    def __call__(self,fname):
 
366
        """Return a runner for the given filename."""
 
367
 
 
368
        if fname.endswith('.py'):
 
369
            runnerClass = PythonRunner
 
370
        elif fname.endswith('.ipy'):
 
371
            runnerClass = IPythonRunner
 
372
        else:
 
373
            raise ValueError('Unknown file type for Runner: %r' % fname)
 
374
 
 
375
        if self.runner is None:
 
376
            return self._makeRunner(runnerClass)
 
377
        else:
 
378
            if runnerClass==self.runnerClass:
 
379
                return self.runner
 
380
            else:
 
381
                e='A runner of type %r can not run file %r' % \
 
382
                   (self.runnerClass,fname)
 
383
                raise ValueError(e)
 
384
 
 
385
 
 
386
# Global usage string, to avoid indentation issues if typed in a function def.
 
387
MAIN_USAGE = """
 
388
%prog [options] file_to_run
 
389
 
 
390
This is an interface to the various interactive runners available in this
 
391
module.  If you want to pass specific options to one of the runners, you need
 
392
to first terminate the main options with a '--', and then provide the runner's
 
393
options.  For example:
 
394
 
 
395
irunner.py --python -- --help
 
396
 
 
397
will pass --help to the python runner.  Similarly,
 
398
 
 
399
irunner.py --ipython -- --interact script.ipy
 
400
 
 
401
will run the script.ipy file under the IPython runner, and then will start to
 
402
interact with IPython at the end of the script (instead of exiting).
 
403
 
 
404
The already implemented runners are listed below; adding one for a new program
 
405
is a trivial task, see the source for examples.
 
406
 
 
407
WARNING: the SAGE runner only works if you manually configure your SAGE copy
 
408
to use 'colors NoColor' in the ipythonrc config file, since currently the
 
409
prompt matching regexp does not identify color sequences.
 
410
"""
 
411
 
 
412
def main():
 
413
    """Run as a command-line script."""
 
414
 
 
415
    parser = optparse.OptionParser(usage=MAIN_USAGE)
 
416
    newopt = parser.add_option
 
417
    parser.set_defaults(mode='ipython')
 
418
    newopt('--ipython',action='store_const',dest='mode',const='ipython',
 
419
           help='IPython interactive runner (default).')
 
420
    newopt('--python',action='store_const',dest='mode',const='python',
 
421
           help='Python interactive runner.')
 
422
    newopt('--sage',action='store_const',dest='mode',const='sage',
 
423
           help='SAGE interactive runner.')
 
424
 
 
425
    opts,args = parser.parse_args()
 
426
    runners = dict(ipython=IPythonRunner,
 
427
                   python=PythonRunner,
 
428
                   sage=SAGERunner)
 
429
 
 
430
    try:
 
431
        ext = os.path.splitext(args[0])[-1]
 
432
    except IndexError:
 
433
        ext = ''
 
434
    modes = {'.ipy':'ipython',
 
435
             '.py':'python',
 
436
             '.sage':'sage'}
 
437
    mode = modes.get(ext,opts.mode)
 
438
    runners[mode]().main(args)
 
439
 
 
440
if __name__ == '__main__':
 
441
    main()