ENH: squeeze keyword to return single col frame as series. #1253

Author: Chang She

pandas/io/parsers.py

@@ -85,6 +85,8 @@ class DateConversionError(Exception):
     Alternative argument name for sep
 encoding : string, default None
     Encoding to use for UTF when reading/writing (ex. 'utf-8')
+squeeze : boolean, default False
+    If the parsed data only contains one column then return a Series
 
 Returns
 -------
@@ -205,7 +207,8 @@ def read_csv(filepath_or_buffer,
              converters=None,
              verbose=False,
              delimiter=None,
-             encoding=None):
+             encoding=None,
+             squeeze=False):
     kwds = locals()
 
     # Alias sep -> delimiter.
@@ -236,7 +239,8 @@ def read_table(filepath_or_buffer,
                converters=None,
                verbose=False,
                delimiter=None,
-               encoding=None):
+               encoding=None,
+               squeeze=False):
     kwds = locals()
 
     # Alias sep -> delimiter.
@@ -271,7 +275,8 @@ def read_fwf(filepath_or_buffer,
              converters=None,
              delimiter=None,
              verbose=False,
-             encoding=None):
+             encoding=None,
+             squeeze=False):
 
     kwds = locals()
 
@@ -372,14 +377,16 @@ class TextParser(object):
         Number of line at bottom of file to skip
     encoding : string, default None
         Encoding to use for UTF when reading/writing (ex. 'utf-8')
+    squeeze : boolean, default False
+        returns Series if only one column
     """
 
     def __init__(self, f, delimiter=None, names=None, header=0,
                  index_col=None, na_values=None, thousands=None,
                  comment=None, parse_dates=False, keep_date_col=False,
                  date_parser=None, dayfirst=False,
                  chunksize=None, skiprows=None, skip_footer=0, converters=None,
-                 verbose=False, encoding=None):
+                 verbose=False, encoding=None, squeeze=False):
         """
         Workhorse function for processing nested list into DataFrame
 
@@ -439,6 +446,7 @@ def __init__(self, f, delimiter=None, names=None, header=0,
         self.index_name = self._get_index_name()
         self._first_chunk = True
 
+        self.squeeze = squeeze
 
     def _make_reader(self, f):
         import csv
@@ -755,7 +763,10 @@ def get_chunk(self, rows=None):
 
         data = _convert_to_ndarrays(data, self.na_values, self.verbose)
 
-        return DataFrame(data=data, columns=columns, index=index)
+        df = DataFrame(data=data, columns=columns, index=index)
+        if self.squeeze and len(df.columns) == 1:
+            return df[df.columns[0]]
+        return df
 
     def _find_line_number(self, exp_len, chunk_len, chunk_i):
         if exp_len is None:

pandas/io/tests/test_parsers.py

@@ -11,11 +11,12 @@
 from numpy import nan
 import numpy as np
 
-from pandas import DataFrame, Index, isnull
+from pandas import DataFrame, Series, Index, isnull
 import pandas.io.parsers as parsers
 from pandas.io.parsers import (read_csv, read_table, read_fwf,
                                ExcelFile, TextParser)
-from pandas.util.testing import assert_almost_equal, assert_frame_equal, network
+from pandas.util.testing import (assert_almost_equal, assert_frame_equal,
+                                 assert_series_equal, network)
 import pandas._tseries as lib
 from pandas.util import py3compat
 from pandas._tseries import Timestamp
@@ -91,6 +92,18 @@ def test_comment_fwf(self):
                       comment='#')
         assert_almost_equal(df.values, expected)
 
+    def test_squeeze(self):
+        data = """\
+a,1
+b,2
+c,3
+"""
+        expected = Series([1,2,3], ['a', 'b', 'c'])
+        result = read_table(StringIO(data), sep=',', index_col=0,
+                            header=None, squeeze=True)
+        self.assert_(isinstance(result, Series))
+        assert_series_equal(result, expected)
+
     def test_multiple_date_col(self):
         # Can use multiple date parsers
         data = """\