0001"""Primary interface layer for terminal control
0002
0003This module provides a semi-portable interface between the various methods of
0004terminal control as the ANSI and win32 methods are quite different.
0005
0006
0007"""
0008
0009import  sys, os, abstract
0010
0011__all__ = ['DISPLAY_CODES', 'COLORS', 'display', 'reset']
0012
0013class Control(object):
0014    """A basic control class.
0015    
0016    Used when terminal support can not be determined.
0017    """
0018
0019    DISPLAY_CODES = ('default', 'bright', 'dim', 'underline', 'blink',
0020                  'reverse', 'hidden')
0021
0022    # Yellow is a bit weird, xterm and rxvt display dark yellow, while linux and 
0023    # Windows display a more brown-ish color. Bright yellow is always yellow.
0024    # order is important here
0025    COLORS = ('black', 'red', 'green', 'yellow', 'blue', 'magenta',
0026               'cyan', 'white')
0027
0028    # setf from curses uses the colors in a different order for 
0029
0030    def display(self, codes=[], fg=None, bg=None):
0031        """Sets the display attributes on the terminal.
0032        
0033        For details on arguments see documentation for color function
0034        in the abstract module.
0035        """
0036        return self._display(*Control.formatcodes(codes, fg, bg))
0037
0038    def _display(self, codes, fg, bg):
0039        pass
0040
0041    def reset(self):
0042        pass
0043
0044    @staticmethod
0045    def formatcodes(codes=[], fg=None, bg=None):
0046        """Makes sure all arguments are valid"""
0047        if isinstance(codes, basestring):
0048            codes=[codes]
0049        else:
0050            codes = list(codes)
0051        for code in codes:
0052            if code not in Control.DISPLAY_CODES:
0053                raise ValueError, ("'%s' not a valid display value" % code)
0054        for color in (fg, bg):
0055            if color != None:
0056                if color not in Control.COLORS:
0057                    raise ValueError, ("'%s' not a valid color" % color)
0058        return [codes, fg, bg]
0059
0060class ANSI(Control):
0061    """ANSI version of terminal control.
0062    
0063    Based on the ANSI X3.64 standard. See http://en.wikipedia.org/wiki/ANSI_X3.64
0064    """
0065
0066    # Don't use these alone
0067    # see http://vt100.net/docs/vt100-ug/chapter3.html
0068    ESCAPE = "\x1b"
0069    CSI = ESCAPE +"["
0070
0071    # see the reset method
0072    RESET_CODE = ESCAPE + "c"
0073
0074    # Display attributes are handled by sending a series of numbers
0075    # along with special characters. For example \x1b[0m will reset
0076    # the display back to the default colors and \x1b[1;4;33;41m will
0077    # result in a very painful underlined bright yellow with red
0078    # background.
0079    CODES={
0080    'default':0,
0081    'bright':1,
0082    'dim':2,
0083    'underline':4,
0084    'blink':5,
0085    'reverse':7,
0086    'hidden':8,
0087    }
0088    FG={
0089    'black':30,
0090    'red':31,
0091    'green':32,
0092    'yellow':33,
0093    'blue':34,
0094    'magenta':35,
0095    'cyan':36,
0096    'white':37,
0097    }
0098    BG={
0099    'black':40,
0100    'red':41,
0101    'green':42,
0102    'yellow':43,
0103    'blue':44,
0104    'magenta':45,
0105    'cyan':46,
0106    'white':47
0107    }
0108
0109    def _display(self, codes, fg, bg):
0110        abstract.stdout.raw_write(ANSI.displaycode(codes, fg, bg))
0111        abstract.stdout.flush()
0112
0113    def reset(self):
0114        """Resets the terminal screen.
0115        
0116        Avoid using this method, it will probably
0117        be moved in the next version.
0118        """
0119        abstract.stdout.raw_write(self.RESET_CODE)
0120
0121    @staticmethod
0122    def displaycode(codes=[], fg=None, bg=None):
0123        """Generates the proper ANSI code"""
0124        codes, fg, bg = Control.formatcodes(codes, fg, bg)
0125        codes = [str(ANSI.CODES[code]) for code in codes]
0126        if fg != None: codes.append(str(ANSI.FG[fg]))
0127        if bg != None: codes.append(str(ANSI.BG[bg]))
0128        return ANSI.CSI + ";".join(codes) + 'm'
0129
0130    @staticmethod
0131    def readcodes(codes):
0132        """Reads a list of codes and generates dict"""
0133        dcodes=[]
0134        fg = bg = None
0135        for code in codes:
0136            code = int(code)
0137            if code in ANSI.FG.values():
0138                fg = code % 10
0139            elif code in ANSI.BG.values():
0140                bg = code % 10
0141            elif code in ANSI.CODES.values():
0142                dcodes.append(code)
0143            else:
0144                pass # drop unhandled values
0145        r = {}
0146        if len(codes): r['codes'] = [Control.DISPLAY_CODES[c] for c in dcodes]
0147        if fg != None: r['fg'] = Control.COLORS[fg]
0148        if bg != None: r['bg'] = Control.COLORS[bg]
0149        return r
0150
0151class Win(Control):
0152    """Windows version of terminal control
0153    
0154    This class should not be used by itself, use either Win32 or WinCTypes
0155    classes that subclasses of this class.
0156    
0157    This class makes extensive use of the Windows API
0158    
0159    
0160    The official documentation for the API is on MSDN:
0161    http://msdn.microsoft.com/library/default.asp?url=/library/en-us/dllproc/base/console_functions.asp
0162    
0163    """
0164
0165    STD_OUTPUT_HANDLE = -11
0166
0167    # These contants are defined in PyWin32
0168    # You can combine the values by doing a bitwise or (|)
0169    # for example FG_BLUE | FG_RED would give magenta (0x05)
0170    #
0171    # these contants are just numbers, It's most useful to think of them in binary,
0172    # but python has no binary number thing, so I'm displaying them in hex
0173    FG_BLUE = 1 << 0
0174    FG_GREEN = 1 << 1
0175    FG_RED = 1 << 2
0176    FG_INTENSITY = 1 << 3
0177    BG_BLUE = 1 << 4
0178    BG_GREEN= 1 << 5
0179    BG_RED  = 1 << 6
0180    BG_INTENSITY = 1 << 7
0181    FG_ALL = FG_BLUE | FG_GREEN | FG_RED
0182    BG_ALL = BG_BLUE | BG_GREEN | BG_RED
0183
0184    # there are also these codes, but according to tcsh's win32/console.c:
0185    # http://cvs.opensolaris.org/source/xref/sfw/usr/src/cmd/tcsh/tcsh-6.12.00/win32/console.c#568
0186    # COMMON_LVB_REVERSE_VIDEO is buggy, so I'm staying away from it. Future
0187    # versions should implement COMMON_LVB_UNDERSCORE.
0188    # COMMON_LVB_REVERSE_VIDEO = 0x4000
0189    # COMMON_LVB_UNDERSCORE      0x8000
0190
0191    default_attributes = None
0192    hidden_output = False
0193    reverse_output = False
0194    reverse_input = False
0195    dim_output = False
0196    real_fg = None
0197
0198    FG = {
0199    'black': 0,
0200    'red': FG_RED,
0201    'green': FG_GREEN,
0202    'yellow': FG_GREEN | FG_RED,
0203    'blue': FG_BLUE,
0204    'magenta': FG_BLUE | FG_RED,
0205    'cyan': FG_BLUE | FG_GREEN,
0206    'white': FG_BLUE | FG_GREEN | FG_RED,
0207    }
0208    BG = {
0209    'black':0,
0210    'red':BG_RED,
0211    'green':BG_GREEN,
0212    'yellow':BG_GREEN | BG_RED,
0213    'blue':BG_BLUE,
0214    'magenta':BG_BLUE | BG_RED,
0215    'cyan':BG_BLUE | BG_GREEN,
0216    'white':BG_BLUE | BG_GREEN | BG_RED,
0217    }
0218
0219    def _set_attributes(self, code):
0220        """ Set console attributes with `code`
0221        
0222        Not implemented here. To be implemented by subclasses.
0223        """
0224        pass
0225
0226    def _split_attributes(self, attrs):
0227        """Spilt attribute code
0228        
0229        Takes an attribute code and returns a tuple containing
0230        foreground (fg), foreground intensity (fgi), background (bg), and
0231        background intensity (bgi)
0232        """
0233        fg = attrs & self.FG_ALL
0234        fgi = attrs & self.FG_INTENSITY
0235        bg = attrs & self.BG_ALL
0236        bgi = attrs & self.BG_INTENSITY
0237        return fg, fgi, bg, bgi
0238
0239    def _cat_attributes(self, fg, fgi, bg, bgi):
0240        return (fg | fgi | bg | bgi)
0241
0242    def _undim(self):
0243        self.dim_output = False
0244        if self.reverse_input:
0245            a = self._get_attributes() & 0x8f
0246            self._set_attributes( (self.real_fg * 0x10) | a)
0247        else:
0248            a = self._get_attributes() & 0xf8
0249            self._set_attributes(self.real_fg | a)
0250
0251    def _display_default(self):
0252        self.hidden_output = False
0253        self.reverse_output = False
0254        self.reverse_input = False
0255        self.dim_output = False
0256        self.real_fg = self.default_attributes & 0x7
0257        self._set_attributes(self.default_attributes)
0258
0259    def _display_bright(self):
0260        self._undim()
0261        return self.FG_INTENSITY
0262
0263    def _display_dim(self):
0264        self.dim_output = True
0265
0266    def _display_reverse(self):
0267        self.reverse_output = True
0268
0269    def _display_hidden(self):
0270        self.hidden_output = True
0271
0272    def _display(self, codes, fg, bg):
0273        color = 0
0274        for c in codes:
0275            try:
0276                f = getattr(self, '_display_' + c)
0277                out = f()
0278                if out: color |= out
0279            except AttributeError:
0280                pass
0281        cfg, cfgi, cbg, cbgi = self._split_attributes(self._get_attributes())
0282        if self.reverse_input:
0283            cfg, cbg = (cbg // 0x10), (cfg * 0x10)
0284            cfgi, cbgi = (cbgi // 0x10), (cfgi * 0x10)
0285        if fg != None:
0286            color |= self.FG[fg]
0287            self.real_fg = self.FG[fg]
0288        else: color |= cfg
0289        if bg != None:
0290            color |= self.BG[bg]
0291        else: color |= cbg
0292        color |= (cfgi | cbgi)
0293        fg, fgi, bg, bgi = self._split_attributes(color)
0294        if self.dim_output:
0295            # intense black
0296            fg = 0
0297            fgi = self.FG_INTENSITY
0298        if self.reverse_output:
0299            fg, bg = (bg // 0x10), (fg * 0x10)
0300            fgi, bgi = (bgi // 0x10), (fgi * 0x10)
0301            self.reverse_input = True
0302        if self.hidden_output:
0303            fg = (bg // 0x10)
0304            fgi = (bgi // 0x10)
0305        self._set_attributes(self._cat_attributes(fg, fgi, bg, bgi))
0306
0307    def reset(self):
0308        """Not yet implemented in the Windows version."""
0309        pass
0310
0311class Win32(Win):
0312    """PyWin32 version of Windows terminal control.
0313    
0314    Uses the PyWin32 Libraries <http://sourceforge.net/projects/pywin32/>.
0315    
0316    ActiveState has good documentation for them:
0317    
0318    Main page:
0319    http://aspn.activestate.com/ASPN/docs/ActivePython/2.4/pywin32/PyWin32.html
0320    Console related objects and methods:
0321    http://aspn.activestate.com/ASPN/docs/ActivePython/2.4/pywin32/PyConsoleScreenBuffer.html
0322    """
0323
0324    def __init__(self):
0325        self._stdout_handle = win32console.GetStdHandle(STD_OUTPUT_HANDLE)
0326        self.default_attributes = self._get_attributes()
0327        self.real_fg = self.default_attributes & 0x7
0328
0329    def _set_attributes(self,code):
0330        self._stdout_handle.SetConsoleTextAttribute(code)
0331
0332    def _get_attributes(self):
0333        return self._stdout_handle.GetConsoleScreenBufferInfo()['Attributes']
0334
0335class WinCTypes(Win):
0336    """CTypes version of Windows terminal control.
0337    
0338    It requires the CTypes libraries <http://sourceforge.net/projects/ctypes/>
0339    
0340    As of Python 2.5, CTypes is included in Python by default. User's of
0341    previous version of Python will have to install it if they what to use
0342    this.
0343    """
0344    def __init__(self):
0345        self._stdout_handle = ctypes.windll.kernel32.GetStdHandle(self.STD_OUTPUT_HANDLE)
0346        self.default_attributes = self._get_attributes()
0347        self.real_fg = self.default_attributes & 0x7
0348
0349    def _get_attributes(self):
0350        # From IPython's winconsole.py, by Alexander Belchenko
0351        # http://projects.scipy.org/ipython/ipython/browser/ipython/trunk/IPython/winconsole.py
0352        import ctypes, struct
0353        csbi = ctypes.create_string_buffer(22)
0354        res = ctypes.windll.kernel32.GetConsoleScreenBufferInfo(self._stdout_handle, csbi)
0355        (bufx, bufy, curx, cury, wattr, left, top, right, bottom, maxx, maxy) =               struct.unpack("hhhhHhhhhhh", csbi.raw)
0357        return wattr
0358
0359    def _set_attributes(self,code):
0360        ctypes.windll.kernel32.SetConsoleTextAttribute(self._stdout_handle, code)
0361
0362def get_wrapper():
0363    """Returns a instance of the Control class (or one of it's sub classes)
0364    
0365    The particular class will depend on the platform, and it should have the proper
0366    methods for the features supported in this module.    
0367    """
0368    ANSI_TERMINALS = ('linux','xterm','rxvt')
0369
0370    # check terminal capabilities
0371    if os.environ.get('TERM') in ANSI_TERMINALS: return ANSI()
0372    elif 'win32' in sys.platform:
0373        try:
0374            import ctypes
0375            return WinCTypes()
0376        except ImportError:
0377            try:
0378                import win32console
0379                return Win32()
0380            except ImportError: return Control()
0381    elif os.environ.get('TERM') == 'cygwin': return ANSI()
0382    else: return Control()
0383
0384
0385wrapper = get_wrapper()
0386
0387# TODO: this old stuff needs some cleanup
0388DISPLAY_CODES = wrapper.DISPLAY_CODES
0389COLORS = wrapper.COLORS
0390display = wrapper.display
0391reset = wrapper.reset
0392
0393# extra for furture versions
0394
0395def _get_capabilities():
0396    """Returns None or a dictionary.
0397    
0398    Don't use this yet unless you are sure you know what your doing, it is 
0399    not well enough tested and is completely missing support for MS Windows.
0400    """
0401
0402    try: import curses
0403    except ImportError: return
0404
0405    # make sure sys.stdout hasn't be hijacked and redirected to a regular
0406    # file
0407    if not sys.stdout.isatty(): return
0408
0409    try: curses.setupterm()
0410    except: return
0411
0412    # bol, eol, bos = beginning/end of line/screen etc
0413    # for info on the values here, see capname's under the terminfo(5) manual
0414    strings = {
0415       'up':'cuu1', 'down':'cud1', 'left':'cub1', 'right':'cuf1',
0416       'bol':'cr', 'bos':'home',
0417       'clear':'clear', 'clear bol':'el1', 'clear eol':'el', 'clear eos':'ed',
0418           'clear line':'dl1',
0419       'bell':'bel',
0420               }
0421
0422    # in the future this should really use sgr or something, but for now
0423    # it pretty much does nothing.
0424    display = {'default':'sgr0','bright':'bold','dim':'dim','reverse':'rev',
0425                       'underline':'smul'}
0426
0427    # dictionary of terminal capabilities. format and display are filled with
0428    # items from format_strings and display_strings. Values are set to None
0429    # if they are not supported.
0430    cap = {
0431    'cols' : None, # if these are None, 75 is a fairly safe fall back to use
0432    'lines' : None,
0433    'fg' : None, # can change foreground
0434    'bg' : None, # can change background
0435    'strings' : dict((string, curses.tigetstr(capname))
0436                    for string, capname in strings.items()),
0437    'display' : dict((string, curses.tigetstr(capname))
0438                     for string, capname in display.items()),
0439    }
0440
0441    for capability in ('cols','lines'):
0442        c = curses.tigetnum(capability)
0443        if c != -1: cap[capability] = c
0444
0445    if curses.tigetstr('setf') != None: cap['fg'] = True
0446    if curses.tigetstr('setb') != None: cap['bg'] = True
0447
0448    for string, capname in format_strings.items():
0449        c = curses.tigetstr(capname)
0450        cap['format'][string] = c