Make all encoding explicit, and blow only raw strings to the terminal.

Tests pass in 2.6.

3 files changed, 103 insertions, 38 deletions

diff --git a/README.rst b/README.rst
index c5aabd7..e95e02a 100644
--- a/README.rst
+++ b/README.rst
@@ -325,6 +325,11 @@ Bugs or suggestions? Visit the `issue tracker`_.
 Version History
 ===============
 
+1.2
+  * Support for Python 3! We need 3.2.3 or greater, because the curses library
+    couldn't decide whether to accept strs or bytes before that
+    (http://bugs.python.org/issue10570).
+
 1.1
   * Added nicely named attributes for colors.
   * Introduced compound formatting.
diff --git a/blessings/__init__.py b/blessings/__init__.py
index 333d68e..6a099ef 100644
--- a/blessings/__init__.py
+++ b/blessings/__init__.py
@@ -8,12 +8,20 @@ except ImportError:
     class IOUnsupportedOperation(Exception):
         """A dummy exception to take the place of Python 3's io.UnsupportedOperation one in Python 2"""
         pass
+import locale
+import os
 from os import isatty, environ
 import struct
 import sys
 from termios import TIOCGWINSZ
 
 
+try:
+    Capability = bytes
+except NameError:
+    Capability = str  # Python 2.5
+
+
 __all__ = ['Terminal']
 __version__ = (1, 1)
 
@@ -38,7 +46,7 @@ class Terminal(object):
         to decide whether to draw progress bars or other frippery.
 
     """
-    def __init__(self, kind=None, stream=None, force_styling=False):
+    def __init__(self, kind=None, stream=None, encoding=None, force_styling=False):
         """Initialize the terminal.
 
         If ``stream`` is not a tty, I will default to returning '' for all
@@ -51,6 +59,10 @@ class Terminal(object):
             the value of the TERM environment variable.
         :arg stream: A file-like object representing the terminal. Defaults to
             the original value of stdout, like ``curses.initscr()`` does.
+        :arg encoding: The encoding to run any Unicode strings through before
+            they're output to the terminal. Terminals take bitstrings, so we've
+            got to pick something. We'll try a pretty fancy cascade of defaults
+            stolen from the standard ``io`` lib if you don't specify.
         :arg force_styling: Whether to force the emission of capabilities, even
             if we don't seem to be in a terminal. This comes in handy if users
             are trying to pipe your output through something like ``less -r``,
@@ -71,6 +83,10 @@ class Terminal(object):
                                  else None)
         except IOUnsupportedOperation:
             stream_descriptor = None
+
+        self.encoding = (self._guess_encoding(stream_descriptor)
+                         if encoding is None else encoding)
+
         self.is_a_tty = stream_descriptor is not None and isatty(stream_descriptor)
         if self.is_a_tty or force_styling:
             # The desciptor to direct terminal initialization sequences to.
@@ -92,10 +108,29 @@ class Terminal(object):
             # that.] At any rate, save redoing the work of _resolve_formatter().
             self._codes = {}
         else:
-            self._codes = NullDict(lambda: NullCallableString(''))
+            self._codes = NullDict(lambda: NullCap('', self.encoding))
 
         self.stream = stream
 
+    def _guess_encoding(self, stream_descriptor):
+        """Return the guessed encoding of the terminal.
+
+        We adapt the algorithm from io.TextIOWrapper, which is what ``print``
+        natively calls in Python 3.
+
+        """
+        encoding = None
+        try:
+            if stream_descriptor is not None:
+                encoding = os.device_encoding(stream_descriptor)
+        except (AttributeError, IOUnsupportedOperation):
+            pass
+        if encoding is None:
+            # getpreferredencoding returns '' in OS X's "Western (NextStep)"
+            # terminal type.
+            encoding = locale.getpreferredencoding() or 'ascii'
+        return encoding
+
     # Sugary names for commonly-used capabilities, intended to help avoid trips
     # to the terminfo man page and comments in your code:
     _sugar = dict(
@@ -188,30 +223,30 @@ class Terminal(object):
         return Location(self, x, y)
 
     def _resolve_formatter(self, attr):
-        """Resolve a sugary or plain capability name, color, or compound formatting function name into a callable string."""
+        """Resolve a sugary or plain capability name, color, or compound formatting function name into a callable capability."""
         if attr in COLORS:
             return self._resolve_color(attr)
         elif attr in COMPOUNDABLES:
             # Bold, underline, or something that takes no parameters
-            return FormattingString(self._resolve_capability(attr), self)
+            return FormattingCap(self._resolve_capability(attr), self)
         else:
             formatters = split_into_formatters(attr)
             if all(f in COMPOUNDABLES for f in formatters):
                 # It's a compound formatter, like "bold_green_on_red". Future
                 # optimization: combine all formatting into a single escape
                 # sequence
-                return FormattingString(''.join(self._resolve_formatter(s)
-                                                for s in formatters),
-                                        self)
+                return FormattingCap(''.join(self._resolve_formatter(s)
+                                             for s in formatters),
+                                     self)
             else:
-                return ParametrizingString(self._resolve_capability(attr))
+                return ParametrizingCap(self._resolve_capability(attr))
 
     def _resolve_capability(self, atom):
         """Return a terminal code for a capname or a sugary name, or ''."""
         return tigetstr(self._sugar.get(atom, atom)) or ''
 
     def _resolve_color(self, color):
-        """Resolve a color like red or on_bright_green into a callable string."""
+        """Resolve a color like red or on_bright_green into a callable capability."""
         # TODO: Does curses automatically exchange red and blue and cyan and
         # yellow when a terminal supports setf/setb rather than setaf/setab?
         # I'll be blasted if I can find any documentation. The following
@@ -222,7 +257,7 @@ class Terminal(object):
         # bright colors at 8-15:
         offset = 8 if 'bright_' in color else 0
         base_color = color.rsplit('_', 1)[-1]
-        return FormattingString(
+        return FormattingCap(
             color_cap(getattr(curses, 'COLOR_' + base_color.upper()) + offset),
             self)
 
@@ -241,8 +276,8 @@ COMPOUNDABLES = (COLORS |
                       'shadow', 'standout', 'subscript', 'superscript']))
 
 
-class ParametrizingString(str):
-    """A string which can be called to parametrize it as a terminal capability"""
+class ParametrizingCap(Capability):
+    """A bytestring which can be called to parametrize it as a terminal capability"""
     def __call__(self, *args):
         try:
             return tparm(self, *args)
@@ -251,7 +286,7 @@ class ParametrizingString(str):
             # running simply `nosetests` (without progressive) on nose-
             # progressive. Perhaps the terminal has gone away between calling
             # tigetstr and calling tparm.
-            return ''
+            return Capability()
         except TypeError:
             # If the first non-int (i.e. incorrect) arg was a string, suggest
             # something intelligent:
@@ -266,11 +301,11 @@ class ParametrizingString(str):
                 raise
 
 
-class FormattingString(str):
-    """A string which can be called upon a piece of text to wrap it in formatting"""
+class FormattingCap(Capability):
+    """A bytestring which can be called upon a piece of text to wrap it in formatting"""
     def __new__(cls, formatting, term):
-        new = str.__new__(cls, formatting)
-        new._term = term
+        new = Capability.__new__(cls, formatting)
+        new._term = term  # TODO: Kill cycle.
         return new
 
     def __call__(self, text):
@@ -283,15 +318,31 @@ class FormattingString(str):
         This should work regardless of whether ``text`` is unicode.
 
         """
+        if isinstance(text, unicode):
+            text = text.encode(self._term.encoding)
         return self + text + self._term.normal
 
 
-class NullCallableString(str):
-    """A callable string that returns '' when called with an int and the arg otherwise."""
+class NullCap(Capability):
+    """A dummy class to stand in for ``FormattingCap`` and ``ParametrizingCap``
+
+    A callable bytestring that returns ``''`` when called with an int and the
+    arg otherwise. We use this when tehre is no tty and so all capabilities are
+    blank.
+
+    """
+    def __new__(cls, cap, encoding):
+        new = Capability.__new__(cls, cap)
+        new._encoding = encoding
+        return new
+
     def __call__(self, arg):
         if isinstance(arg, int):
-            return ''
-        return arg
+            return Capability()
+        elif isinstance(arg, unicode):
+            return arg.encode(self._encoding)
+        else:
+            return arg
 
 
 class NullDict(defaultdict):
diff --git a/blessings/tests.py b/blessings/tests.py
index 62de8fd..d8b5fcd 100644
--- a/blessings/tests.py
+++ b/blessings/tests.py
@@ -10,6 +10,12 @@ from nose.tools import eq_
 # This tests that __all__ is correct, since we use below everything that should
 # be imported:
 from blessings import *
+from blessings import Capability
+
+
+def bytes_eq(bytes1, bytes2):
+    """Make sure ``bytes1`` equals ``bytes2``, the latter of which gets cast to something bytes-like, depending on Python version."""
+    eq_(bytes1, Capability(bytes2))
 
 
 def test_capability():
@@ -28,13 +34,14 @@ def test_capability():
 def test_capability_without_tty():
     """Assert capability templates are '' when stream is not a tty."""
     t = Terminal(stream=StringIO())
-    eq_(t.save, '')
-    eq_(t.red, '')
+    eq_(t.save, Capability(''))
+    eq_(t.red, Capability(''))
 
 
 def test_capability_with_forced_tty():
+    """If we force styling, capabilities had better not (generally) be empty."""
     t = Terminal(stream=StringIO(), force_styling=True)
-    assert t.save != ''
+    assert len(t.save) > 0
 
 
 def test_parametrization():
@@ -63,7 +70,7 @@ def test_location():
 
     eq_(t.stream.getvalue(), tigetstr('sc') +
                              tparm(tigetstr('cup'), 4, 3) +
-                             'hi' +
+                             'hi' +  # TODO: Encode with Terminal's encoding.
                              tigetstr('rc'))
 
 def test_horizontal_location():
@@ -75,7 +82,7 @@ def test_horizontal_location():
 
 
 def test_null_fileno():
-    """Make sure ``Terinal`` works when ``fileno`` is ``None``.
+    """Make sure ``Terminal`` works when ``fileno`` is ``None``.
 
     This simulates piping output to another program.
 
@@ -83,7 +90,7 @@ def test_null_fileno():
     out = stream=StringIO()
     out.fileno = None
     t = Terminal(stream=out)
-    eq_(t.save, '')
+    bytes_eq(t.save, '')
 
 
 def test_mnemonic_colors():
@@ -91,22 +98,23 @@ def test_mnemonic_colors():
     # Avoid testing red, blue, yellow, and cyan, since they might someday
     # chance depending on terminal type.
     t = Terminal()
-    eq_(t.white, '\x1b[37m')
-    eq_(t.green, '\x1b[32m')  # Make sure it's different than white.
-    eq_(t.on_black, '\x1b[40m')
-    eq_(t.on_green, '\x1b[42m')
-    eq_(t.bright_black, '\x1b[90m')
-    eq_(t.bright_green, '\x1b[92m')
-    eq_(t.on_bright_black, '\x1b[100m')
-    eq_(t.on_bright_green, '\x1b[102m')
+    bytes_eq(t.white, '\x1b[37m')
+    bytes_eq(t.green, '\x1b[32m')  # Make sure it's different than white.
+    bytes_eq(t.on_black, '\x1b[40m')
+    bytes_eq(t.on_green, '\x1b[42m')
+    bytes_eq(t.bright_black, '\x1b[90m')
+    bytes_eq(t.bright_green, '\x1b[92m')
+    bytes_eq(t.on_bright_black, '\x1b[100m')
+    bytes_eq(t.on_bright_green, '\x1b[102m')
 
 
 def test_formatting_functions():
     """Test crazy-ass formatting wrappers, both simple and compound."""
-    t = Terminal()
+    t = Terminal(encoding='utf-8')
     eq_(t.bold('hi'), t.bold + 'hi' + t.normal)
     eq_(t.green('hi'), t.green + 'hi' + t.normal)
-    eq_(t.bold_green(u'boö'), t.bold + t.green + u'boö' + t.normal)  # unicode
+    # Test encoding of unicodes:
+    eq_(t.bold_green(u'boö'), t.bold + t.green + u'boö'.encode('utf-8') + t.normal)
     eq_(t.bold_underline_green_on_red('boo'),
         t.bold + t.underline + t.green + t.on_red + 'boo' + t.normal)
     # Don't spell things like this:
@@ -119,7 +127,8 @@ def test_formatting_functions_without_tty():
     t = Terminal(stream=StringIO())
     eq_(t.bold('hi'), 'hi')
     eq_(t.green('hi'), 'hi')
-    eq_(t.bold_green(u'boö'), u'boö')  # unicode
+    # Test encoding of unicodes:
+    eq_(t.bold_green(u'boö'), u'boö'.encode('utf-8'))  # unicode
     eq_(t.bold_underline_green_on_red('boo'), 'boo')
     eq_(t.on_bright_red_bold_bright_green_underline('meh'), 'meh')