gsh/stdin.py

   1 # This program is free software; you can redistribute it and/or modify
   2 # it under the terms of the GNU General Public License as published by
   3 # the Free Software Foundation; either version 2 of the License, or
   4 # (at your option) any later version.
   5 #
   6 # This program is distributed in the hope that it will be useful,
   7 # but WITHOUT ANY WARRANTY; without even the implied warranty of
   8 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   9 # GNU Library General Public License for more details.
  10 #
  11 # You should have received a copy of the GNU General Public License
  12 # along with this program; if not, write to the Free Software
  13 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  14 #
  15 # See the COPYING file for license information.
  16 #
  17 # Copyright (c) 2006 Guillaume Chazarain <guichaz@yahoo.fr>
  18
  19 import asyncore
  20 import atexit
  21 import errno
  22 import fcntl
  23 import os
  24 import readline # Just to say we want to use it with raw_input
  25 import sys
  26 from threading import Thread, Event, Lock
  27
  28 from gsh import remote_dispatcher
  29 from gsh.console import set_blocking_stdin, console_output
  30
  31 # Handling of stdin is certainly the most complex part of gsh
  32
  33 def restore_streams_flags_at_exit():
  34     """We play we fcntl flags, so we make sure to restore them on exit"""
  35     get_flags = lambda fd: (fd, fcntl.fcntl(fd, fcntl.F_GETFL, 0))
  36     flags = map(get_flags, range(3))
  37     set_flags = lambda (fd, flags): fcntl.fcntl(fd, fcntl.F_SETFL, flags)
  38     atexit.register(map, set_flags, flags)
  39
  40 class stdin_dispatcher(asyncore.file_dispatcher):
  41     """The stdin reader in the main thread => no fancy editing"""
  42     def __init__(self):
  43         asyncore.file_dispatcher.__init__(self, 0)
  44         self.is_readable = True
  45         set_blocking_stdin(True)
  46
  47     def readable(self):
  48         """We set it to be readable only when the stdin thread is not in
  49         raw_input()"""
  50         return self.is_readable
  51
  52     def writable(self):
  53         """We don't write to stdin"""
  54         return False
  55
  56     def handle_close(self):
  57         """Ctrl-D was received but the remote processes were not ready"""
  58         remote_dispatcher.dispatch_termination_to_all()
  59
  60     def handle_read(self):
  61         """Some data can be read on stdin"""
  62         while True:
  63             try:
  64                 set_blocking_stdin(False)
  65                 try:
  66                     data = self.recv(4096)
  67                 finally:
  68                     set_blocking_stdin(True)
  69             except OSError, e:
  70                 if e.errno == errno.EAGAIN:
  71                     # End of available data
  72                     break
  73                 else:
  74                     raise
  75             else:
  76                 if data:
  77                     # Handle the just read data
  78                     the_stdin_thread.input_buffer.add(data)
  79                     process_input_buffer()
  80                 else:
  81                     # Closed?
  82                     self.is_readable = False
  83                     break
  84
  85 class input_buffer(object):
  86     """The shared input buffer between the main thread and the stdin thread"""
  87     def __init__(self):
  88         self.lock = Lock()
  89         self.buf = ''
  90
  91     def add(self, data):
  92         """Add data to the buffer"""
  93         self.lock.acquire()
  94         try:
  95             self.buf += data
  96         finally:
  97             self.lock.release()
  98
  99
 100     def get(self):
 101         """Get the content of the buffer"""
 102         self.lock.acquire()
 103         try:
 104             data = self.buf
 105             if data:
 106                 self.buf = ''
 107                 return data
 108         finally:
 109             self.lock.release()
 110
 111 def process_input_buffer():
 112     """Send the content of the input buffer to all remote processes, this must
 113     be called in the main thread"""
 114     data = the_stdin_thread.input_buffer.get()
 115     if not data:
 116         return
 117     for r in remote_dispatcher.all_instances():
 118         try:
 119             r.dispatch_write(data)
 120         except Exception, msg:
 121             console_output('%s for %s, disconnecting\n' % (msg, r.display_name),
 122                            output=sys.stderr)
 123             r.disconnect()
 124         else:
 125             r.log(('<== ', data))
 126             if r.enabled and r.state == remote_dispatcher.STATE_IDLE:
 127                 r.change_state(remote_dispatcher.STATE_EXPECTING_NEXT_LINE)
 128
 129 # The stdin thread uses a pipe to communicate with the main thread, which is
 130 # most of the time waiting in the select() loop.
 131 # Pipe character protocol:
 132 # s: entering in raw_input, the main loop should not read stdin
 133 # e: leaving raw_input, the main loop can read stdin
 134 # q: Ctrl-D was pressed, exiting
 135 # d: there is new data to send
 136
 137 class pipe_notification_reader(asyncore.file_dispatcher):
 138     """The pipe reader in the main thread"""
 139     def __init__(self):
 140         asyncore.file_dispatcher.__init__(self, the_stdin_thread.pipe_read)
 141
 142     def _do(self, c):
 143         if c in ('s', 'e'):
 144             the_stdin_dispatcher.is_readable = c == 'e'
 145         elif c == 'q':
 146             remote_dispatcher.dispatch_termination_to_all()
 147         elif c == 'd':
 148             process_input_buffer()
 149         else:
 150             raise Exception, 'Unknown code: %s' % (c)
 151
 152     def handle_read(self):
 153         """Handle all the available character commands in the pipe"""
 154         while True:
 155             try:
 156                 c = self.recv(1)
 157             except OSError, e:
 158                 ok = e.errno == errno.EAGAIN
 159                 assert ok
 160                 return
 161             else:
 162                 self._do(c)
 163
 164 # All the words that have been typed in gsh. Used by the completion mechanism.
 165 history_words = set()
 166
 167 def complete(text, state):
 168     """On tab press, return the next possible completion"""
 169     l = len(text)
 170     matches = [w for w in history_words if len(w) > l and w.startswith(text)]
 171     if state <= len(matches):
 172         return matches[state]
 173
 174 class stdin_thread(Thread):
 175     """The stdin thread, used to call raw_input()"""
 176     def __init__(self):
 177         Thread.__init__(self, name='stdin thread')
 178
 179     @staticmethod
 180     def activate(interactive):
 181         """Activate the thread at initialization time"""
 182         the_stdin_thread.ready_event = Event()
 183         the_stdin_thread.input_buffer = input_buffer()
 184         if interactive:
 185             the_stdin_thread.interrupted_event = Event()
 186             the_stdin_thread.pipe_read, the_stdin_thread.pipe_write = os.pipe()
 187             the_stdin_thread.wants_control_shell = False
 188             the_stdin_thread.setDaemon(True)
 189             the_stdin_thread.start()
 190             pipe_notification_reader()
 191         else:
 192             the_stdin_thread.ready_event.set()
 193
 194     def run(self):
 195         while True:
 196             self.ready_event.wait()
 197             # The remote processes are ready, the thread can call raw_input
 198             self.interrupted_event.clear()
 199             console_output('\r')
 200             set_blocking_stdin(True)
 201             try:
 202                 try:
 203                     os.write(self.pipe_write, 's')
 204                     nr = remote_dispatcher.count_completed_processes()[0]
 205                     readline.set_completer(complete)
 206                     readline.parse_and_bind('tab: complete')
 207                     readline.set_completer_delims(' \t\n')
 208                     cmd = raw_input('gsh (%d)> ' % (nr))
 209                     if self.wants_control_shell:
 210                         # This seems to be needed if Ctrl-C is hit when some
 211                         # text is in the line buffer
 212                         raise EOFError
 213                     if len(history_words) < 10000:
 214                         for word in cmd.split():
 215                             if len(word) > 1:
 216                                 history_words.add(word)
 217                 finally:
 218                     set_blocking_stdin(False)
 219                     os.write(self.pipe_write, 'e')
 220             except EOFError:
 221                 if self.wants_control_shell:
 222                     self.ready_event.clear()
 223                     # Ok, we are no more in raw_input(), tell it to the
 224                     # main thread
 225                     self.interrupted_event.set()
 226                 else:
 227                     os.write(self.pipe_write, 'q')
 228                     return
 229             else:
 230                 self.ready_event.clear()
 231                 self.input_buffer.add(cmd + '\n')
 232                 os.write(self.pipe_write, 'd')
 233
 234 the_stdin_thread = stdin_thread()
 235 the_stdin_dispatcher = stdin_dispatcher()