Allow unicode formulas on Py2 in limited circumstances.

njsmith · njsmith · commit 1a19ab796b70 · 2015-10-26T22:54:57.000-07:00
Apparently our disallowing unicode formula strings on Py2 was being rather annoying for people using `from __future__ import unicode_literals`. Start allowing them in limited circumstances. Fixes pydatagh-53.
diff --git a/doc/changes.rst b/doc/changes.rst
@@ -6,6 +6,14 @@ Changes
 v0.4.1
 ------
 
+New features:
+
+* On Python 2, accept ``unicode`` strings containing only ASCII
+  characters as valid formula descriptions in
+  the high-level formula API (:func:`dmatrix` and friends). This is
+  intended as a convenience for people using Python 2 with ``from
+  __future__ import unicode_literals``. (See :ref:`py2-versus-py3`.)
+
 Bug fixes:
 
 * Accept ``long`` as a valid integer type in the new
diff --git a/doc/py2-versus-py3.rst b/doc/py2-versus-py3.rst
@@ -1,3 +1,5 @@
+.. _py2-versus-py3:
+
 Python 2 versus Python 3
 ========================
 
@@ -6,11 +8,11 @@ Python 2 versus Python 3
 The biggest difference between Python 2 and Python 3 is in their
 string handling, and this is particularly relevant to Patsy since
 it parses user input. We follow a simple rule: input to Patsy
-should always be of type `str`. That means that on Python 2, you
+should always be of type ``str``. That means that on Python 2, you
 should pass byte-strings (not unicode), and on Python 3, you should
 pass unicode strings (not byte-strings). Similarly, when Patsy
 passes text back (e.g. :attr:`DesignInfo.column_names`), it's always
-in the form of a `str`.
+in the form of a ``str``.
 
 In addition to this being the most convenient for users (you never
 need to use any b"weird" u"prefixes" when writing a formula string),
@@ -20,3 +22,14 @@ byte-strings, and that's the only form of input accepted by the
 :mod:`tokenize` module. On the other hand, Python 3's tokenizer and
 parser use unicode, and since Patsy processes Python code, it has
 to follow suit.
+
+There is one exception to this rule: on Python 2, as a convenience for
+those using ``from __future__ import unicode_literals``, the
+high-level API functions :func:`dmatrix`, :func:`dmatrices`,
+:func:`incr_dbuilders`, and :func:`incr_dbuilder` do accept
+``unicode`` strings -- BUT these unicode string objects are still
+required to contain only ASCII characters; if they contain any
+non-ASCII characters then an error will be raised. If you really need
+non-ASCII in your formulas, then you should consider upgrading to
+Python 3. Low-level APIs like :meth:`ModelDesc.from_formula` continue
+to insist on ``str`` objects only.
diff --git a/patsy/highlevel.py b/patsy/highlevel.py
@@ -13,6 +13,7 @@
 #   ModelDesign doesn't work -- need to work with the builder set
 #   want to be able to return either a matrix or a pandas dataframe
 
+import six
 import numpy as np
 from patsy import PatsyError
 from patsy.design_info import DesignMatrix, DesignInfo
@@ -45,6 +46,18 @@ def _try_incr_builders(formula_like, data_iter_maker, eval_env,
             raise PatsyError("bad value from %r.__patsy_get_model_desc__"
                                 % (formula_like,))
         # fallthrough
+    if not six.PY3 and isinstance(formula_like, unicode):
+        # Included for the convenience of people who are using py2 with
+        # __future__.unicode_literals.
+        try:
+            formula_like = formula_like.encode("ascii")
+        except UnicodeEncodeError:
+            raise PatsyError(
+                "On Python 2, formula strings must be either 'str' objects, "
+                "or else 'unicode' objects containing only ascii "
+                "characters. You passed a unicode string with non-ascii "
+                "characters. I'm afraid you'll have to either switch to "
+                "ascii-only, or else upgrade to Python 3.")
     if isinstance(formula_like, str):
         formula_like = ModelDesc.from_formula(formula_like)
         # fallthrough
diff --git a/patsy/test_highlevel.py b/patsy/test_highlevel.py
@@ -6,6 +6,7 @@
 
 import sys
 import __future__
+import six
 import numpy as np
 from nose.tools import assert_raises
 from patsy import PatsyError
@@ -74,7 +75,7 @@ def t(formula_like, data, depth,
         depth += 1
     def data_iter_maker():
         return iter([data])
-    if (isinstance(formula_like, (str, ModelDesc, DesignInfo))
+    if (isinstance(formula_like, six.string_types + (ModelDesc, DesignInfo))
         or (isinstance(formula_like, tuple)
             and isinstance(formula_like[0], DesignInfo))
         or hasattr(formula_like, "__patsy_get_model_desc__")):
@@ -258,7 +259,21 @@ def __patsy_get_model_desc__(self, data):
     t("x + y", {"y": [1, 2], "x": [3, 4]}, 0,
       True,
       [[1, 3, 1], [1, 4, 2]], ["Intercept", "x", "y"])
-    
+
+    # unicode objects on py2 (must be ascii only)
+    if not six.PY3:
+        # ascii is fine
+        t(unicode("y ~ x"),
+          {"y": [1, 2], "x": [3, 4]}, 0,
+          True,
+          [[1, 3], [1, 4]], ["Intercept", "x"],
+          [[1], [2]], ["y"])
+        # non-ascii is not (even if this would be valid on py3 with its less
+        # restrict variable naming rules)
+        eacute = "\xc3\xa9".decode("utf-8")
+        assert isinstance(eacute, unicode)
+        assert_raises(PatsyError, dmatrix, eacute, data={eacute: [1, 2]})
+
     # ModelDesc
     desc = ModelDesc([], [Term([LookupFactor("x")])])
     t(desc, {"x": [1.5, 2.5, 3.5]}, 0,
