Call 'str' on index columns that have an object or categorical dtype (we were already doing this on non-index columns).

Author: Tim Shawver

qgrid/grid.py

@@ -748,16 +748,36 @@ def _update_table(self,
 
         self._row_count = len(self._df.index)
 
-        if type(df.index) == pd.core.index.MultiIndex:
-            self._multi_index = True
-        else:
-            self._multi_index = False
-
         if update_columns:
             self._string_columns = list(df.select_dtypes(
                 include=[np.dtype('O'), 'category']
             ).columns.values)
 
+            def should_be_stringified(col_series):
+                return col_series.dtype == np.dtype('O') or \
+                       hasattr(col_series, 'cat')
+
+            if type(df.index) == pd.core.index.MultiIndex:
+                self._multi_index = True
+                for idx, cur_level in enumerate(df.index.levels):
+                    if cur_level.name:
+                        col_name = cur_level.name
+                    else:
+                        col_name = 'level_%s' % idx
+                    self._primary_key.append(col_name)
+                    if should_be_stringified(cur_level):
+                        self._string_columns.append(col_name)
+            else:
+                self._multi_index = False
+                if df.index.name:
+                    col_name = df.index.name
+                else:
+                    col_name = 'index'
+                self._primary_key = [col_name]
+
+                if should_be_stringified(df.index):
+                    self._string_columns.append(col_name)
+
         # call map(str) for all columns identified as string columns, in
         # case any are not strings already
         for col_name in self._string_columns:
@@ -1102,7 +1122,7 @@ def _get_col_series_from_df(self, col_name, df):
         if col_name in self._primary_key:
             if len(self._primary_key) > 1:
                 key_index = self._primary_key.index(col_name)
-                return df.index.get_level_values(level=key_index)
+                return df.index.levels[key_index]
             else:
                 return df.index
         else:

qgrid/tests/test_grid.py

@@ -76,6 +76,7 @@ def test_edit_date():
 
 def test_edit_multi_index_df():
     df_multi = create_multi_index_df()
+    df_multi.index.set_names('first', level=0, inplace=True)
     view = QgridWidget(df=df_multi)
     old_val = df_multi.loc[('bar', 'two'), 1]
 
@@ -493,7 +494,7 @@ def __init__(self, obj):
 
 
 def test_object_dtype():
-    df = pd.DataFrame({'a': my_object_vals})
+    df = pd.DataFrame({'a': my_object_vals}, index=my_object_vals)
     widget = QgridWidget(df=df)
     grid_data = json.loads(widget._df_json)['data']
 
@@ -520,6 +521,20 @@ def test_object_dtype():
     assert not isinstance(grid_data[0]['a'], dict)
     assert not isinstance(grid_data[1]['a'], dict)
 
+    assert not isinstance(grid_data[0]['index'], dict)
+    assert not isinstance(grid_data[1]['index'], dict)
+
+
+def test_index_categorical():
+    df = pd.DataFrame({'foo': np.random.randn(3), 'future_index': [22, 13, 87]})
+    df['future_index'] = df['future_index'].astype('category')
+    df = df.set_index('future_index')
+    widget = QgridWidget(df=df)
+    grid_data = json.loads(widget._df_json)['data']
+
+    assert not isinstance(grid_data[0]['future_index'], dict)
+    assert not isinstance(grid_data[1]['future_index'], dict)
+
 
 def test_object_dtype_categorical():
     cat_series = pd.Series(