Ticket #2157: _win32stdio.patch

File _win32stdio.patch, 7.4 KB (added by synapsis, 8 years ago)

patch for _win32stdio.py, version 2

  • _win32stdio.py

     
    1 # -*- test-case-name: twisted.test.test_process.ProcessTestCase.testStdio -*- 
     1# -*- test-case-name: twisted.test.test_process.ProcessTestCase.testStdio,twisted.test.test_conio.StdIOTestCase -*- 
    22 
     3import os 
     4import errno 
     5import msvcrt 
     6 
     7import pywintypes 
    38import win32api 
    4 import os, msvcrt 
     9import win32console 
    510 
    611from zope.interface import implements 
    712 
    813from twisted.internet.interfaces import IHalfCloseableProtocol, ITransport, IAddress 
    914from twisted.internet.interfaces import IConsumer, IPushProducer 
     15from twisted.internet import main, abstract 
    1016 
    11 from twisted.internet import _pollingfile, main 
     17import _pollingfile 
    1218 
     19 
     20 
     21# _pollingfile support 
     22# XXX check me 
     23class _PollableReadConsole(_pollingfile._PollableReader): 
     24    def __init__(self, channel, receivedCallback, lostCallback): 
     25        _pollingfile._PollableReader.__init__(self, channel.handle, 
     26                                              receivedCallback, lostCallback)  
     27        self.channel = channel 
     28 
     29    def checkWork(self): 
     30        try: 
     31            data = self.channel.read(2000) # 2000 = 80 x 25 
     32        except IOError, ioe: 
     33            assert ioe.args[0] == errno.EAGAIN 
     34            return 0 
     35        except win32console.error: 
     36            # stdin or stdout closed? 
     37            self.cleanup() 
     38            return 0 
     39 
     40        self.receivedCallback(data) 
     41        return len(data) 
     42 
     43 
     44class _PollableWriteConsole(_pollingfile._PollableWriter): 
     45    def __init__(self, channel, lostCallback): 
     46        _pollingfile._PollableWriter.__init__(self, channel.handle, lostCallback) 
     47 
     48        self.channel = channel 
     49 
     50    def checkWork(self): 
     51        numBytesWritten = 0 
     52        if not self.outQueue: 
     53            if self.disconnecting: 
     54                self.writeConnectionLost() 
     55                return 0 
     56            try: 
     57                self.channel.write('') 
     58            except pywintypes.error: 
     59                self.writeConnectionLost() 
     60                return numBytesWritten 
     61 
     62        while self.outQueue: 
     63            data = self.outQueue.pop(0) 
     64            try: 
     65                # XXX as far as I know,  
     66                # nBytesWritten is always equal to len(data) 
     67                nBytesWritten = self.channel.write(data) 
     68            except win32console.error: 
     69                self.writeConnectionLost() 
     70                break 
     71            else: 
     72                numBytesWritten += nBytesWritten 
     73                if len(data) > nBytesWritten: 
     74                    self.outQueue.insert(0, data[nBytesWritten:]) 
     75                    break 
     76        else: 
     77            resumed = self.bufferEmpty() 
     78            if not resumed and self.disconnecting: 
     79                self.writeConnectionLost() 
     80 
     81        return numBytesWritten 
     82 
     83 
     84 
     85# support for win32eventreactor 
     86# XXX check me 
     87class ConsoleReader(abstract.FileDescriptor): 
     88    def __init__(self, channel, receivedCallback, lostCallback, reactor=None): 
     89        """win32eventreactor is assumed. 
     90        """ 
     91         
     92        if not reactor: 
     93            from twisted.internet import reactor 
     94 
     95        self.channel = channel 
     96        self.receivedCallback = receivedCallback 
     97        self.lostCallback = lostCallback 
     98        self.reactor = reactor 
     99 
     100        self.reactor.addEvent(self.channel.handle, self, "doRead") 
     101 
     102 
     103    def doRead(self): 
     104        try: 
     105            data = self.channel.read(2000) # 2000 = 80 x 25 
     106        except IOError, ioe: 
     107            assert ioe.args[0] == errno.EAGAIN 
     108            return 0 
     109        except win32console.error: 
     110            # stdin or stdout closed? 
     111            self.cleanup() 
     112            return 0 
     113 
     114        self.receivedCallback(data) 
     115        return len(data) 
     116 
     117    def cleanup(self): 
     118        self.reactor.removeEvent(self.channel.handle) 
     119        self.lostCallback() 
     120 
     121    def close(self): 
     122        try: 
     123            win32api.CloseHandle(self.channel.handle) 
     124        except pywintypes.error: 
     125            # You can't close std handles...? 
     126            pass 
     127 
     128        self.cleanup() 
     129 
     130    def stopProducing(self): 
     131        self.close() 
     132 
     133    def pauseProducing(self): 
     134        self.reactor.removeEvent(self.channel.handle) 
     135 
     136    def resumeProducing(self): 
     137        self.reactor.addEvent(self.channel.handle, self, "doRead") 
     138 
     139 
    13140class Win32PipeAddress(object): 
    14141    implements(IAddress) 
    15142 
     143class Win32ConsoleAddress(object): 
     144    implements(IAddress) 
     145 
     146 
    16147class StandardIO(_pollingfile._PollingTimer): 
    17148 
    18149    implements(ITransport, 
     
    28159 
    29160        Also, put it stdin/stdout/stderr into binary mode. 
    30161        """ 
     162 
    31163        from twisted.internet import reactor 
    32  
     164        from twisted.internet.win32eventreactor import Win32Reactor 
     165         
    33166        for stdfd in range(0, 1, 2): 
    34167            msvcrt.setmode(stdfd, os.O_BINARY) 
    35168 
    36169        _pollingfile._PollingTimer.__init__(self, reactor) 
    37170        self.proto = proto 
     171         
     172        # Check for win32eventreactor 
     173        win32Enabled = reactor.__class__ is Win32Reactor 
    38174 
    39         hstdin = win32api.GetStdHandle(win32api.STD_INPUT_HANDLE) 
    40         hstdout = win32api.GetStdHandle(win32api.STD_OUTPUT_HANDLE) 
     175        # Check if we are connected to a console. 
     176        # If this is the case, connect to the console, else connect to 
     177        # anonymous pipes. 
     178        if os.isatty(0): 
     179            import conio 
     180             
     181            if win32Enabled: 
     182                self.stdin = ConsoleReader( 
     183                    conio.stdin, self.dataReceived, self.readConnectionLost 
     184                    ) 
     185            else: 
     186                self.stdin = _PollableReadConsole( 
     187                    conio.stdin, self.dataReceived, self.readConnectionLost 
     188                    ) 
     189        else: 
     190            hstdin = win32api.GetStdHandle(win32api.STD_INPUT_HANDLE) 
     191            self.stdin = _pollingfile._PollableReadPipe( 
     192                hstdin, self.dataReceived, self.readConnectionLost 
     193                ) 
    41194 
    42         self.stdin = _pollingfile._PollableReadPipe( 
    43             hstdin, self.dataReceived, self.readConnectionLost) 
     195        if os.isatty(1): 
     196            import conio 
     197             
     198            self.stdout = _PollableWriteConsole( 
     199                conio.stdout, self.writeConnectionLost 
     200                ) 
     201        else: 
     202            hstdout = win32api.GetStdHandle(win32api.STD_OUTPUT_HANDLE) 
     203            self.stdout = _pollingfile._PollableWritePipe( 
     204                hstdout, self.writeConnectionLost 
     205                ) 
     206             
     207        if not (os.isatty(0) and win32Enabled): 
     208            self._addPollableResource(self.stdin) 
    44209 
    45         self.stdout = _pollingfile._PollableWritePipe( 
    46             hstdout, self.writeConnectionLost) 
    47  
    48         self._addPollableResource(self.stdin) 
    49210        self._addPollableResource(self.stdout) 
    50211 
    51212        self.proto.makeConnection(self) 
     
    86247        self.stdout.close() 
    87248 
    88249    def getPeer(self): 
    89         return Win32PipeAddress() 
     250        if os.isatty(0) and os.isatty(1): 
     251            return Win32ConsoleAddress() 
     252        else: 
     253            return Win32PipeAddress() 
    90254 
    91255    def getHost(self): 
    92         return Win32PipeAddress() 
     256        if os.isatty(0) and os.isatty(1): 
     257            return Win32ConsoleAddress() 
     258        else: 
     259            return Win32PipeAddress() 
    93260 
     261 
    94262    # IConsumer 
    95263 
    96264    def registerProducer(self, producer, streaming): 
     
    113281 
    114282    def resumeProducing(self): 
    115283        self.stdin.resumeProducing() 
    116