BUG: Inherit iterator parsers from BaseIterator

closes pandas-dev#12153

Author: toobaz

pandas/io/common.py

@@ -408,7 +408,7 @@ def UnicodeReader(f, dialect=csv.excel, encoding="utf-8", **kwds):
     def UnicodeWriter(f, dialect=csv.excel, encoding="utf-8", **kwds):
         return csv.writer(f, dialect=dialect, **kwds)
 else:
-    class UnicodeReader:
+    class UnicodeReader(BaseIterator):
 
         """
         A CSV reader which will iterate over lines in the CSV file "f",
@@ -422,16 +422,10 @@ def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds):
             f = UTF8Recoder(f, encoding)
             self.reader = csv.reader(f, dialect=dialect, **kwds)
 
-        def next(self):
+        def __next__(self):
             row = next(self.reader)
             return [compat.text_type(s, "utf-8") for s in row]
 
-        # python 3 iterator
-        __next__ = next
-
-        def __iter__(self):  # pragma: no cover
-            return self
-
     class UnicodeWriter:
 
         """

pandas/io/parsers.py

@@ -19,7 +19,8 @@
 from pandas.core.config import get_option
 from pandas.io.date_converters import generic_parser
 from pandas.io.common import (get_filepath_or_buffer, _validate_header_arg,
-                              _get_handle, UnicodeReader, UTF8Recoder)
+                              _get_handle, UnicodeReader, UTF8Recoder,
+                              BaseIterator)
 from pandas.tseries import tools
 
 from pandas.util.decorators import Appender
@@ -545,7 +546,7 @@ def read_fwf(filepath_or_buffer, colspecs='infer', widths=None, **kwds):
 ])
 
 
-class TextFileReader(object):
+class TextFileReader(BaseIterator):
     """
 
     Passed dialect overrides any of the related parser options
@@ -724,15 +725,8 @@ def _clean_options(self, options, engine):
 
         return result, engine
 
-    def __iter__(self):
-        try:
-            if self.chunksize:
-                while True:
-                    yield self.read(self.chunksize)
-            else:
-                yield self.read()
-        except StopIteration:
-            pass
+    def __next__(self):
+        return self.get_chunk()
 
     def _make_engine(self, engine='c'):
         if engine == 'c':
@@ -2363,7 +2357,7 @@ def _concat_date_cols(date_cols):
     return rs
 
 
-class FixedWidthReader(object):
+class FixedWidthReader(BaseIterator):
     """
     A reader of fixed-width lines.
     """
@@ -2417,7 +2411,7 @@ def detect_colspecs(self, n=100):
         edges = np.where((mask ^ shifted) == 1)[0]
         return list(zip(edges[::2], edges[1::2]))
 
-    def next(self):
+    def __next__(self):
         if self.buffer is not None:
             try:
                 line = next(self.buffer)
@@ -2430,9 +2424,6 @@ def next(self):
         return [line[fromm:to].strip(self.delimiter)
                 for (fromm, to) in self.colspecs]
 
-    # Iterator protocol in Python 3 uses __next__()
-    __next__ = next
-
 
 class FixedWidthFieldParser(PythonParser):
     """

pandas/io/sas.py

@@ -10,7 +10,7 @@
 
 from datetime import datetime
 import pandas as pd
-from pandas.io.common import get_filepath_or_buffer
+from pandas.io.common import get_filepath_or_buffer, BaseIterator
 from pandas import compat
 import struct
 import numpy as np
@@ -242,7 +242,7 @@ def _parse_float_vec(vec):
     return ieee
 
 
-class XportReader(object):
+class XportReader(BaseIterator):
     __doc__ = _xport_reader_doc
 
     def __init__(self, filepath_or_buffer, index=None, encoding='ISO-8859-1',
@@ -369,15 +369,8 @@ def _read_header(self):
         dtype = np.dtype(dtypel)
         self._dtype = dtype
 
-    def __iter__(self):
-        try:
-            if self._chunksize:
-                while True:
-                    yield self.read(self._chunksize)
-            else:
-                yield self.read()
-        except StopIteration:
-            pass
+    def __next__(self):
+        return self.read(nrows=self._chunksize or 1)
 
     def _record_count(self):
         """

pandas/io/stata.py

@@ -25,7 +25,7 @@
 from pandas.util.decorators import Appender
 import pandas as pd
 import pandas.core.common as com
-from pandas.io.common import get_filepath_or_buffer
+from pandas.io.common import get_filepath_or_buffer, BaseIterator
 from pandas.lib import max_len_string_array, infer_dtype
 from pandas.tslib import NaT, Timestamp
 
@@ -907,7 +907,7 @@ def _decode_bytes(self, str, errors=None):
             return str
 
 
-class StataReader(StataParser):
+class StataReader(StataParser, BaseIterator):
     __doc__ = _stata_reader_doc
 
     def __init__(self, path_or_buf, convert_dates=True,
@@ -1377,15 +1377,8 @@ def data(self, **kwargs):
 
         return self.read(None, **kwargs)
 
-    def __iter__(self):
-        try:
-            if self._chunksize:
-                while True:
-                    yield self.read(self._chunksize)
-            else:
-                yield self.read()
-        except StopIteration:
-            pass
+    def __next__(self):
+        return self.read(nrows=self._chunksize or 1)
 
     def get_chunk(self, size=None):
         """