ENH: more intelligent inference about index_col for Excel files, test coverage for PR #735

Author: wesm

pandas/core/frame.py

@@ -1,3 +1,5 @@
+from __future__ import with_statement
+
 """
 DataFrame
 ---------
@@ -853,12 +855,12 @@ def _helper_csvexcel(self, writer, na_rep=None, cols=None, header=True,
                         index_label = []
                         for i, name in enumerate(self.index.names):
                             if name is None:
-                                name = 'level_%d' % i
+                                name = '' # 'level_%d' % i
                             index_label.append(name)
                     else:
                         index_label = self.index.name
                         if index_label is None:
-                            index_label = ['index']
+                            index_label = ['']
                         else:
                             index_label = [index_label]
                 elif not isinstance(index_label, (list, tuple, np.ndarray)):
@@ -917,8 +919,8 @@ def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
         ----------
         path : string
             File path
-        nanRep : string, default ''
-            Missing data rep'n
+        na_rep : string, default ''
+            Missing data representation
         cols : sequence, optional
             Columns to write
         header : boolean, default True
@@ -936,18 +938,17 @@ def to_csv(self, path, sep=",", na_rep='', cols=None, header=True,
             a string representing the encoding to use if the contents are
             non-ascii, for python versions prior to 3
         """
-        f = open(path, mode)
-        csvout = csv.writer(f, lineterminator='\n', delimiter=sep)
-
         if nanRep is not None:  # pragma: no cover
             import warnings
             warnings.warn("nanRep is deprecated, use na_rep",
                           FutureWarning)
             na_rep = nanRep
 
-        self._helper_csvexcel(csvout, na_rep=na_rep, cols=cols, header=header,
-                         index=index, index_label=index_label, encoding=encoding)
-        f.close()
+        with open(path, mode) as f:
+            csvout = csv.writer(f, lineterminator='\n', delimiter=sep)
+            self._helper_csvexcel(csvout, na_rep=na_rep, cols=cols,
+                                  header=header, index=index,
+                                  index_label=index_label, encoding=encoding)
 
     def to_excel(self, excel_writer, sheet_name = 'sheet1', na_rep='', cols=None, header=True,
                  index=True, index_label=None):

pandas/io/parsers.py

@@ -259,7 +259,7 @@ def _make_reader(self, f):
         if sep is None or len(sep) == 1:
             sniff_sep = True
             # default dialect
-            dia = csv.excel
+            dia = csv.excel()
             if sep is not None:
                 sniff_sep = False
                 dia.delimiter = sep
@@ -477,10 +477,7 @@ def get_chunk(self, rows=None):
         for col, f in self.converters.iteritems():
             if isinstance(col, int) and col not in self.columns:
                 col = self.columns[col]
-            result = lib.map_infer(data[col], f)
-            if issubclass(result.dtype.type, (basestring, unicode)):
-                result = result.astype('O')
-            data[col] = result
+            data[col] = lib.map_infer(data[col], f)
 
         data = _convert_to_ndarrays(data, self.na_values, self.verbose)
 
@@ -620,8 +617,14 @@ def _parse_xlsx(self, sheetname, header=0, skiprows=None, index_col=None,
               chunksize=None):
         sheet = self.book.get_sheet_by_name(name=sheetname)
         data = []
-        for row in sheet.iter_rows(): # it brings a new method: iter_rows()
+
+        # it brings a new method: iter_rows()
+        for row in sheet.iter_rows():
             data.append([cell.internal_value for cell in row])
+
+        if header is not None:
+            data[header] = _trim_excel_header(data[header])
+
         parser = TextParser(data, header=header, index_col=index_col,
                             na_values=na_values,
                             parse_dates=parse_dates,
@@ -630,7 +633,7 @@ def _parse_xlsx(self, sheetname, header=0, skiprows=None, index_col=None,
                             chunksize=chunksize)
 
         return parser.get_chunk()
-        
+
     def _parse_xls(self, sheetname, header=0, skiprows=None, index_col=None,
               parse_dates=False, date_parser=None, na_values=None,
               chunksize=None):
@@ -654,6 +657,9 @@ def _parse_xls(self, sheetname, header=0, skiprows=None, index_col=None,
                 row.append(value)
             data.append(row)
 
+        if header is not None:
+            data[header] = _trim_excel_header(data[header])
+
         parser = TextParser(data, header=header, index_col=index_col,
                             na_values=na_values,
                             parse_dates=parse_dates,
@@ -663,9 +669,15 @@ def _parse_xls(self, sheetname, header=0, skiprows=None, index_col=None,
 
         return parser.get_chunk()
 
+def _trim_excel_header(row):
+    # trim header row so auto-index inference works
+    while len(row) > 0 and row[0] == '':
+        row = row[1:]
+    return row
+
 class ExcelWriter(object):
     """
-    Class for writing DataFrame objects into excel sheets, uses xlwt for xls, 
+    Class for writing DataFrame objects into excel sheets, uses xlwt for xls,
     openpyxl for xlsx.  See DataFrame.to_excel for typical usage.
 
     Parameters
@@ -701,14 +713,15 @@ def writerow(self, row, sheet_name=None):
         Parameters
         ----------
         row : list
-            Row of data to save to Excel sheet 
+            Row of data to save to Excel sheet
         sheet_name : string, default None
             Name of Excel sheet, if None, then use self.cur_sheet
         """
         if sheet_name is None:
             sheet_name = self.cur_sheet
-        if sheet_name is None:
-            raise Exception('Must pass explicit sheet_name or set cur_sheet property')
+        if sheet_name is None:  # pragma: no cover
+            raise Exception('Must pass explicit sheet_name or set '
+                            'cur_sheet property')
         if self.use_xlsx:
             self._writerow_xlsx(row, sheet_name)
         else:
@@ -720,13 +733,13 @@ def _writerow_xls(self, row, sheet_name):
         else:
             sheet = self.book.add_sheet(sheet_name)
             row_idx = 0
-        sheetrow = sheet.row(row_idx) 
+        sheetrow = sheet.row(row_idx)
         for i, val in enumerate(row):
             if isinstance(val, (datetime.datetime, datetime.date)):
                 if isinstance(val, datetime.datetime):
-                    sheetrow.write(i,val,self.fm_datetime)
+                    sheetrow.write(i,val, self.fm_datetime)
                 else:
-                    sheetrow.write(i,val,self.fm_date)
+                    sheetrow.write(i,val, self.fm_date)
             elif isinstance(val, np.int64):
                 sheetrow.write(i,int(val))
             else:
@@ -744,6 +757,7 @@ def _writerow_xlsx(self, row, sheet_name):
             sheet.title = sheet_name
             row_idx = 0
 
-        sheet.append([int(val) if isinstance(val, np.int64) else val for val in row])
+        sheet.append([int(val) if isinstance(val, np.int64) else val
+                      for val in row])
         row_idx += 1
         self.sheets[sheet_name] = (sheet, row_idx)

pandas/io/tests/test_parsers.py

@@ -287,6 +287,27 @@ def test_sniff_delimiter(self):
         data2 = read_csv(StringIO(text), index_col=0, delimiter='|')
         assert_frame_equal(data, data2)
 
+        text = """ignore this
+ignore this too
+index|A|B|C
+foo|1|2|3
+bar|4|5|6
+baz|7|8|9
+"""
+        data3 = read_csv(StringIO(text), index_col=0, sep=None, skiprows=2)
+        assert_frame_equal(data, data3)
+
+        text = u"""ignore this
+ignore this too
+index|A|B|C
+foo|1|2|3
+bar|4|5|6
+baz|7|8|9
+""".encode('utf-8')
+        data4 = read_csv(StringIO(text), index_col=0, sep=None, skiprows=2,
+                         encoding='utf-8')
+        assert_frame_equal(data, data4)
+
     def test_read_nrows(self):
         df = read_csv(StringIO(self.data1), nrows=3)
         expected = read_csv(StringIO(self.data1))[:3]

pandas/tests/test_frame.py

@@ -2265,6 +2265,11 @@ def test_to_csv_unicode(self):
         df.to_csv(path, encoding='UTF-8')
         df2 = read_csv(path, index_col=0, encoding='UTF-8')
         assert_frame_equal(df, df2)
+
+        df.to_csv(path, encoding='UTF-8', index=False)
+        df2 = read_csv(path, index_col=None, encoding='UTF-8')
+        assert_frame_equal(df, df2)
+
         os.remove(path)
 
     def test_to_excel_from_excel(self):
@@ -2281,40 +2286,40 @@ def test_to_excel_from_excel(self):
             # test roundtrip
             self.frame.to_excel(path,'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0)
+            recons = reader.parse('test1', index_col=0)
             assert_frame_equal(self.frame, recons)
-            
+
             self.frame.to_excel(path,'test1', index=False)
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=None)
+            recons = reader.parse('test1', index_col=None)
             recons.index = self.frame.index
             assert_frame_equal(self.frame, recons)
 
             self.frame.to_excel(path,'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0,skiprows=[1])
+            recons = reader.parse('test1', index_col=0, skiprows=[1])
             assert_frame_equal(self.frame.ix[1:], recons)
 
             self.frame.to_excel(path,'test1',na_rep='NA')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0,na_values=['NA'])
+            recons = reader.parse('test1', index_col=0, na_values=['NA'])
             assert_frame_equal(self.frame, recons)
-            
+
             self.mixed_frame.to_excel(path,'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0)
+            recons = reader.parse('test1', index_col=0)
             assert_frame_equal(self.mixed_frame, recons)
 
-            self.tsframe.to_excel(path,'test1')
+            self.tsframe.to_excel(path, 'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0)
+            recons = reader.parse('test1')
             assert_frame_equal(self.tsframe, recons)
 
             #Test np.int64
             frame = DataFrame(np.random.randn(10,2))
             frame.to_excel(path,'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1',index_col=0)
+            recons = reader.parse('test1')
             assert_frame_equal(frame, recons)
 
             # Test writing to separate sheets
@@ -2330,14 +2335,25 @@ def test_to_excel_from_excel(self):
 
             os.remove(path)
 
+        # datetime.date, not sure what to test here exactly
+        path = '__tmp__.xls'
+        tsf = self.tsframe.copy()
+        tsf.index = [x.date() for x in self.tsframe.index]
+        tsf.to_excel(path, 'test1')
+        reader = ExcelFile(path)
+        recons = reader.parse('test1')
+        assert_frame_equal(self.tsframe, recons)
+        os.remove(path)
+
     def test_to_excel_multiindex(self):
         for ext in ['xls', 'xlsx']:
             path = '__tmp__.' + ext
 
             frame = self.frame
             old_index = frame.index
             arrays = np.arange(len(old_index)*2).reshape(2,-1)
-            new_index = MultiIndex.from_arrays(arrays, names=['first', 'second'])
+            new_index = MultiIndex.from_arrays(arrays,
+                                               names=['first', 'second'])
             frame.index = new_index
             frame.to_excel(path, 'test1', header=False)
             frame.to_excel(path, 'test1', cols=['A', 'B'])
@@ -2361,11 +2377,22 @@ def test_to_excel_multiindex(self):
             recons = reader.parse('test1', index_col=[0,1])
             assert_frame_equal(tsframe, recons)
 
-            # do not load index
+            # infer index
             tsframe.to_excel(path, 'test1')
             reader = ExcelFile(path)
-            recons = reader.parse('test1', index_col=None)
-            np.testing.assert_equal(len(recons.columns), len(tsframe.columns) + 2)
+            recons = reader.parse('test1')
+            assert_frame_equal(tsframe, recons)
+
+            # no index
+            tsframe.index.names = ['first', 'second']
+            tsframe.to_excel(path, 'test1')
+            reader = ExcelFile(path)
+            recons = reader.parse('test1')
+            assert_almost_equal(tsframe.values,
+                                recons.ix[:, tsframe.columns].values)
+            self.assertEqual(len(tsframe.columns) + 2, len(recons.columns))
+
+            tsframe.index.names = [None, None]
 
             # no index
             tsframe.to_excel(path, 'test1', index=False)
@@ -2374,6 +2401,10 @@ def test_to_excel_multiindex(self):
             assert_almost_equal(recons.values, self.tsframe.values)
             self.tsframe.index = old_index # needed if setUP becomes classmethod
 
+            # write a big DataFrame
+            df = DataFrame(np.random.randn(1005, 1))
+            df.to_excel(path, 'test1')
+
             os.remove(path)
 
     def test_info(self):

pandas/tools/tests/test_merge.py

@@ -543,6 +543,12 @@ def test_handle_join_key_pass_array(self):
         self.assert_(merged['key'].notnull().all())
         self.assert_(merged2['key'].notnull().all())
 
+        left = DataFrame({'value' : range(5)}, columns=['value', 'key'])
+        right = DataFrame({'rvalue' : range(6)})
+        lkey = np.array([1, 1, 2, 2, 3])
+        rkey = np.array([1, 1, 2, 3, 4, 5])
+
+
 class TestMergeMulti(unittest.TestCase):
 
     def setUp(self):