| mbligh | 56f1fbb | 2006-10-01 15:10:56 +0000 | [diff] [blame] | 1 | __author__ = """Copyright Martin J. Bligh, Andy Whitcroft, 2005, 2006""" |
| 2 | |
| 3 | import sys, os |
| 4 | |
| 5 | class fd_stack: |
| 6 | """a stack of fd redirects |
| 7 | |
| 8 | Redirects cause existing fd's to be pushed on the stack; restore() |
| 9 | causes the current set of redirects to be popped, restoring the previous |
| 10 | filehandle destinations. |
| 11 | |
| 12 | Note that we need to redirect both the sys.stdout type descriptor |
| 13 | (which print, etc use) and the low level OS numbered descriptor |
| 14 | which os.system() etc use. |
| 15 | """ |
| 16 | |
| 17 | def __init__(self, fd, filehandle): |
| 18 | self.fd = fd # eg 1 |
| 19 | self.filehandle = filehandle # eg sys.stdout |
| 20 | self.stack = [(fd, filehandle)] |
| 21 | |
| 22 | |
| 23 | def update_handle(self, new): |
| 24 | if (self.filehandle == sys.stdout): |
| 25 | sys.stdout = new |
| 26 | if (self.filehandle == sys.stderr): |
| 27 | sys.stderr = new |
| 28 | self.filehandle = new |
| 29 | |
| 30 | def redirect(self, filename): |
| 31 | """Redirect output to the specified file |
| 32 | |
| 33 | Overwrites the previous contents, if any. |
| 34 | """ |
| 35 | self.filehandle.flush() |
| 36 | fdcopy = os.dup(self.fd) |
| 37 | self.stack.append( (fdcopy, self.filehandle, 0) ) |
| 38 | # self.filehandle = file(filename, 'w') |
| 39 | if (os.path.isfile(filename)): |
| 40 | newfd = os.open(filename, os.O_WRONLY) |
| 41 | else: |
| 42 | newfd = os.open(filename, os.O_WRONLY | os.O_CREAT) |
| 43 | os.dup2(newfd, self.fd) |
| 44 | os.close(newfd) |
| 45 | self.update_handle(os.fdopen(self.fd, 'w')) |
| 46 | |
| 47 | |
| 48 | def tee_redirect(self, filename): |
| 49 | """Tee output to the specified file |
| 50 | |
| 51 | Overwrites the previous contents, if any. |
| 52 | """ |
| 53 | self.filehandle.flush() |
| 54 | #print_to_tty("tee_redirect to " + filename) |
| 55 | #where_art_thy_filehandles() |
| 56 | fdcopy = os.dup(self.fd) |
| 57 | r, w = os.pipe() |
| 58 | pid = os.fork() |
| 59 | if pid: # parent |
| 60 | os.close(r) |
| 61 | os.dup2(w, self.fd) |
| 62 | os.close(w) |
| 63 | else: # child |
| 64 | os.close(w) |
| 65 | os.dup2(r, 0) |
| 66 | os.dup2(fdcopy, 1) |
| 67 | os.close(r) |
| 68 | os.close(fdcopy) |
| 69 | os.execlp('tee', 'tee', filename) |
| 70 | self.stack.append( (fdcopy, self.filehandle, pid) ) |
| 71 | self.update_handle(os.fdopen(self.fd, 'w')) |
| 72 | #where_art_thy_filehandles() |
| 73 | #print_to_tty("done tee_redirect to " + filename) |
| 74 | |
| 75 | |
| 76 | def restore(self): |
| 77 | """unredirect one level""" |
| 78 | self.filehandle.flush() |
| 79 | # print_to_tty("ENTERING RESTORE %d" % self.fd) |
| 80 | # where_art_thy_filehandles() |
| 81 | (old_fd, old_filehandle, pid) = self.stack.pop() |
| 82 | # print_to_tty("old_fd %d" % old_fd) |
| 83 | # print_to_tty("self.fd %d" % self.fd) |
| 84 | self.filehandle.close() # seems to close old_fd as well. |
| 85 | if pid: |
| 86 | os.waitpid(pid, 0) |
| 87 | # where_art_thy_filehandles() |
| 88 | os.dup2(old_fd, self.fd) |
| 89 | # print_to_tty("CLOSING FD %d" % old_fd) |
| 90 | os.close(old_fd) |
| 91 | # where_art_thy_filehandles() |
| 92 | self.update_handle(old_filehandle) |
| 93 | # where_art_thy_filehandles() |
| 94 | # print_to_tty("EXIT RESTORE %d" % self.fd) |