handle pandas.Categorical-valued factors as categorical data

njsmith · njsmith · commit f9e8b2036991 · 2013-02-09T18:02:32.000-08:00
diff --git a/patsy/build.py b/patsy/build.py
@@ -441,6 +441,9 @@ def _examine_factor_types(factors, factor_states, data_iter_maker):
             break
         for factor in list(examine_needed):
             value = factor.eval(factor_states[factor], data)
+            if have_pandas and isinstance(value, pandas.Categorical):
+                value = Categorical.from_pandas_categorical(value)
+                # fall through into the next 'if':
             if isinstance(value, Categorical):
                 postprocessor = CategoricalTransform(levels=value.levels)
                 prefinished_postprocessors[factor] = postprocessor
diff --git a/patsy/categorical.py b/patsy/categorical.py
@@ -9,8 +9,8 @@
 from patsy import PatsyError
 from patsy.state import stateful_transform
 from patsy.util import (SortAnythingKey,
-                           have_pandas, asarray_or_pandas,
-                           pandas_friendly_reshape)
+                        have_pandas, asarray_or_pandas,
+                        pandas_friendly_reshape)
 
 if have_pandas:
     import pandas
@@ -39,6 +39,11 @@ def __init__(self, int_array, levels, contrast=None):
         self.levels = tuple(levels)
         self.contrast = contrast
 
+    @classmethod
+    def from_pandas_categorical(cls, pandas_categorical):
+        return Categorical(pandas_categorical.labels,
+                           pandas_categorical.levels)
+
     @classmethod
     def from_sequence(cls, sequence, levels=None, **kwargs):
         """from_sequence(sequence, levels=None, contrast=None)
@@ -166,8 +171,8 @@ def test_Categorical():
 class CategoricalTransform(object):
     """C(data, contrast=None, levels=None)
 
-    Converts some `data` into :class:`Categorical` form. (It is also used
-    called implicitly any time a formula contains a bare categorical factor.)
+    Converts some `data` into :class:`Categorical` form. (It is also called
+    implicitly any time a formula contains a bare categorical factor.)
 
     This is used in two cases:
 
@@ -209,6 +214,9 @@ def memorize_finish(self):
 
     def transform(self, data, contrast=None, levels=None):
         kwargs = {"contrast": contrast}
+        if have_pandas and isinstance(data, pandas.Categorical):
+            data = Categorical.from_pandas_categorical(data)
+            # fall through to the next 'if':
         if isinstance(data, Categorical):
             if levels is not None and data.levels != levels:
                 raise PatsyError("changing levels of categorical data "
@@ -286,6 +294,13 @@ def test_C_pandas():
         assert np.array_equal(cat3.int_array.index, [10, 20, 30])
         assert cat3.contrast == "asdf"
 
+def test_categorical_from_pandas_categorical():
+    if have_pandas:
+        pandas_categorical = pandas.Categorical.from_array(["a", "b", "a"])
+        c = Categorical.from_pandas_categorical(pandas_categorical)
+        assert np.array_equal(c.int_array, [0, 1, 0])
+        assert c.levels == ("a", "b")
+
 def test_categorical_non_strings():
     cat = C([1, "foo", ("a", "b")])
     assert set(cat.levels) == set([1, "foo", ("a", "b")])
diff --git a/patsy/test_build.py b/patsy/test_build.py
@@ -436,14 +436,19 @@ def iter_maker():
 def test_categorical():
     data_strings = {"a": ["a1", "a2", "a1"]}
     data_categ = {"a": C(["a2", "a1", "a2"])}
+    datas = [data_strings, data_categ]
+    if have_pandas:
+        data_pandas = {"a": pandas.Categorical.from_array(["a1", "a2", "a2"])}
+        datas.append(data_pandas)
     def t(data1, data2):
         def iter_maker():
             yield data1
         builders = design_matrix_builders([make_termlist(["a"])],
                                           iter_maker)
         build_design_matrices(builders, data2)
-    t(data_strings, data_categ)
-    t(data_categ, data_strings)
+    for data1 in datas:
+        for data2 in datas:
+            t(data1, data2)
 
 def test_contrast():
     from patsy.contrasts import ContrastMatrix, Sum
