convert unicode user categories -> utf8 under python2

3 files changed, 25 insertions, 4 deletions

diff --git a/CHANGES b/CHANGES
index 0c58167..8d4ad15 100644
--- a/CHANGES
+++ b/CHANGES
@@ -9,6 +9,11 @@ Release History
 
     Minor bugfix release
     
+    * Various :meth:`~passlib.context.CryptContext` methods
+      would incorrectly raise :exc:`TypeError` if passed a :class:`!unicode`
+      user category under Python 2; for compatibility
+      they will now be treated the same as the equivalent ``utf-8`` :class:`bytes`.
+      
     * *bugfix*: FreeBSD 8.3 added native support for SHA512-Crypt,
       updated unittests and documentation accordingly (:issue:`35`).
       
diff --git a/passlib/context.py b/passlib/context.py
index dca68ae..12ad5e5 100644
--- a/passlib/context.py
+++ b/passlib/context.py
@@ -21,7 +21,7 @@ from passlib.registry import get_crypt_handler, _validate_handler_name
 from passlib.utils import rng, tick, to_bytes, \
                           to_unicode, splitcomma
 from passlib.utils.compat import bytes, iteritems, num_types, \
-                                 PY3, PY_MIN_32, unicode, SafeConfigParser, \
+                                 PY2, PY3, PY_MIN_32, unicode, SafeConfigParser, \
                                  NativeStringIO, BytesIO, base_string_types
 #pkg
 #local
@@ -2087,9 +2087,12 @@ class CryptContext(object):
 
         # type check
         if category is not None and not isinstance(category, str):
-            raise ExpectedTypeError(category, "str|None", "category")
+            if PY2 and isinstance(category, unicode):
+                # for compatibility with unicode-centric py2 apps
+                return self._get_record(scheme, category.encode("utf-8"))
+            raise ExpectedTypeError(category, "str or None", "category")
         if scheme is not None and not isinstance(scheme, str):
-            raise ExpectedTypeError(scheme, "str|None", "scheme")
+            raise ExpectedTypeError(scheme, "str or None", "scheme")
 
         # if scheme=None, use category's default scheme, and cache result.
         if not scheme:
diff --git a/passlib/tests/test_context.py b/passlib/tests/test_context.py
index 90709d1..a9e52fd 100644
--- a/passlib/tests/test_context.py
+++ b/passlib/tests/test_context.py
@@ -22,7 +22,7 @@ from passlib import hash
 from passlib.context import CryptContext, LazyCryptContext
 from passlib.exc import PasslibConfigWarning
 from passlib.utils import tick, to_bytes, to_unicode
-from passlib.utils.compat import irange, u, unicode, str_to_uascii
+from passlib.utils.compat import irange, u, unicode, str_to_uascii, PY2
 import passlib.utils.handlers as uh
 from passlib.tests.utils import TestCase, catch_warnings, set_file, TICK_RESOLUTION, quicksleep
 from passlib.registry import (register_crypt_handler_path,
@@ -659,6 +659,11 @@ sha512_crypt__min_rounds = 45000
         self.assertEqual(ctx.handler(category="staff"), hash.sha256_crypt)
         self.assertEqual(ctx.handler(category="admin"), hash.md5_crypt)
 
+        # test unicode category strings are accepted under py2
+        if PY2:
+            self.assertEqual(ctx.handler(category=u("staff")), hash.sha256_crypt)
+            self.assertEqual(ctx.handler(category=u("admin")), hash.md5_crypt)
+
     def test_33_options(self):
         "test internal _get_record_options() method"
         def options(ctx, scheme, category=None):
@@ -836,6 +841,14 @@ sha512_crypt__min_rounds = 45000
         # border cases
         #--------------------------------------------------------------
 
+        # test unicode category strings are accepted under py2
+        # this tests basic _get_record() used by encrypt/genhash/verify.
+        # we have to omit scheme=xxx so codepath is tested fully
+        if PY2:
+            c2 = cc.copy(default="phpass")
+            self.assertTrue(c2.genconfig(category=u("admin")).startswith("$P$5"))
+            self.assertTrue(c2.genconfig(category=u("staff")).startswith("$H$5"))
+
         # throws error without schemes
         self.assertRaises(KeyError, CryptContext().genconfig)
         self.assertRaises(KeyError, CryptContext().genconfig, scheme='md5_crypt')