Merge pull request #159 from jtkiley/master

davidastephens · davidastephens · commit c88216105143 · 2016-01-29T21:54:37.000-08:00
Added retrieval of EDGAR daily indices.
diff --git a/docs/source/remote_data.rst b/docs/source/remote_data.rst
@@ -404,8 +404,21 @@ EDGAR Index
 
 Company filing index from EDGAR (SEC).
 
+The daily indices get large quickly (i.e. the set of daily indices from 1994
+to 2015 is 1.5GB), and the FTP server will close the connection past some
+downloading threshold . In testing, pulling one year at a time works well.
+If the FTP server starts refusing your connections, you should be able to
+reconnect after waiting a few minutes.
+
+
 .. ipython:: python
 
     import pandas_datareader.data as web
     ed = web.DataReader('full', 'edgar-index')
     ed[:5]
+
+.. ipython:: python
+
+    import pandas_datareader.data as web
+    ed = web.DataReader('daily', 'edgar-index', '1998-05-18', '1998-05-18')
+    ed[:5]
diff --git a/docs/source/whatsnew/v0.2.2.txt b/docs/source/whatsnew/v0.2.2.txt
@@ -20,6 +20,7 @@ New features
 
 - ``DataReader`` now supports dividend only pulls from Yahoo! Finance, see :ref:`here<remote_data.yahoo>` (:issue:`138`).
 - ``DataReader`` now supports SEC EDGAR full (current quarter) index retrieval, see :ref:`here<remote_data.edgar>` (:issue:`143`).
+- ``DataReader`` now supports SEC EDGAR daily (back to 7/1/1994) index retrieval, see :ref:`here<remote_data.edgar>` (:issue:`147`).
 
 .. _whatsnew_022.api_breaking:
 
diff --git a/pandas_datareader/compat/__init__.py b/pandas_datareader/compat/__init__.py
@@ -0,0 +1 @@
+from io import BytesIO
diff --git a/pandas_datareader/data.py b/pandas_datareader/data.py
@@ -94,7 +94,8 @@ def DataReader(name, data_source=None, start=None, end=None,
     ff = DataReader("F-F_ST_Reversal_Factor", "famafrench")
 
     # Data from EDGAR index
-    ed = DataReader("master", "edgar-index")
+    ed = DataReader("full", "edgar-index")
+    ed2 = DataReader("daily", "edgar-index")
     """
     if data_source == "yahoo":
         return YahooDailyReader(symbols=name, start=start, end=end,
diff --git a/pandas_datareader/edgar.py b/pandas_datareader/edgar.py
@@ -1,13 +1,39 @@
+import re
+import datetime as dt
+from ftplib import FTP
+import gzip
+
 from pandas import read_csv
+from pandas import DataFrame
+from pandas import to_datetime
 from pandas.io.common import ZipFile
 from pandas.compat import StringIO
-from pandas.compat import BytesIO
+from pandas.core.common import is_number
 
 from pandas_datareader.base import _BaseReader
+from pandas_datareader.compat import BytesIO
+from pandas_datareader._utils import RemoteDataError
+
 
+_URL_FULL = 'edgar/full-index/master.zip'
+_URL_DAILY = 'ftp://ftp.sec.gov/'
+_SEC_FTP = 'ftp.sec.gov'
 
-_URL_FULL = 'ftp://ftp.sec.gov/edgar/full-index/master.zip'
 _COLUMNS = ['cik', 'company_name', 'form_type', 'date_filed', 'filename']
+_DIVIDER = re.compile('--------------')
+_EDGAR = 'edgar/'
+_EDGAR_DAILY = 'edgar/daily-index'
+_EDGAR_RE = re.compile(_EDGAR)
+_EDGAR_MIN_DATE = dt.datetime(1994, 7, 1)
+_ZIP_RE = re.compile('\.zip$')
+_GZ_RE = re.compile('\.gz$')
+
+_MLSD_VALUES_RE = re.compile('modify=(?P<modify>.*?);.*'
+                             'type=(?P<type>.*?);.*'
+                             '; (?P<name>.*)$')
+_FILENAME_DATE_RE = re.compile('\w*?\.(\d*)\.idx')
+_FILENAME_MASTER_RE = re.compile('master\.\d*\.idx')
+_EDGAR_MAX_6_DIGIT_DATE = dt.datetime(1998, 5, 15)
 
 
 class EdgarIndexReader(_BaseReader):
@@ -17,27 +43,202 @@ class EdgarIndexReader(_BaseReader):
     Returns
     -------
     edgar_index : pandas.DataFrame.
-        DataFrame of EDGAR master index.
+        DataFrame of EDGAR index.
     """
 
     @property
     def url(self):
-        return _URL_FULL
-
-    def _read_zipfile(self, url):
+        if self.symbols == 'full':
+            return _URL_FULL
+        elif self.symbols == 'daily':
+            return _URL_DAILY
+        else:
+            return _URL_FULL  # Should probably raise or use full unless daily.
 
-        zipf = BytesIO(self._get_response(url).content)
+    def _read_zipfile(self, ftppath):
 
+        zipf = BytesIO()
+        try:
+            self._sec_ftp_session.retrbinary('RETR ' + ftppath, zipf.write)
+        except EOFError:
+            raise RemoteDataError('FTP server has closed the connection.')
+        zipf.seek(0)
         with ZipFile(zipf, 'r') as zf:
             data = zf.open(zf.namelist()[0]).read().decode()
 
-        return data
+        return StringIO(data)
+
+    def _read_gzfile(self, ftppath):
+
+        zipf = BytesIO()
+        try:
+            self._sec_ftp_session.retrbinary('RETR ' + ftppath, zipf.write)
+        except EOFError:
+            raise RemoteDataError('FTP server has closed the connection.')
+        zipf.seek(0)
+        zf = gzip.GzipFile(fileobj=zipf, mode='rb')
+        try:
+            data = zf.read().decode('iso-8859-1')
+        finally:
+            zf.close()
+
+        return StringIO(data)
 
-    def _read_one_data(self, url, params):
+    def _read_one_data(self, ftppath, params):
 
-        index_file = StringIO(self._read_zipfile(url))
+        if re.search(_ZIP_RE, ftppath) is not None:
+            index_file = self._read_zipfile(ftppath)
+        elif re.search(_GZ_RE, ftppath) is not None:
+            index_file = self._read_gzfile(ftppath)
+        else:
+            index_file = StringIO()
+            index_list = []
+            try:
+                self._sec_ftp_session.retrlines('RETR ' + ftppath,
+                                                index_list.append)
+            except EOFError:
+                raise RemoteDataError('FTP server has closed the connection.')
 
+            for line in index_list:
+                index_file.write(line + '\n')
+            index_file.seek(0)
+
+        index_file = self._remove_header(index_file)
         index = read_csv(index_file, delimiter='|', header=None,
-                         index_col=False, skiprows=10, names=_COLUMNS,
+                         index_col=False, names=_COLUMNS,
                          low_memory=False)
+        index['filename'] = index['filename'].map(self._fix_old_file_paths)
         return index
+
+    def _read_daily_data(self, url, params):
+        doc_index = DataFrame()
+        file_index = self._get_dir_lists()
+        for idx_entry in file_index:
+            if self._check_idx(idx_entry):
+                daily_idx_path = (idx_entry['path'] + '/' + idx_entry['name'])
+                daily_idx = self._read_one_data(daily_idx_path, params)
+                doc_index = doc_index.append(daily_idx)
+        return doc_index
+
+    def _check_idx(self, idx_entry):
+        if re.match(_FILENAME_MASTER_RE, idx_entry['name']):
+            if idx_entry['date'] is not None:
+                if (self.start <= idx_entry['date'] <= self.end):
+                    return True
+        else:
+            return False
+
+    def _remove_header(self, data):
+        header = True
+        cleaned_datafile = StringIO()
+        for line in data:
+            if header is False:
+                cleaned_datafile.write(line + '\n')
+            elif re.search(_DIVIDER, line) is not None:
+                header = False
+
+        cleaned_datafile.seek(0)
+        return cleaned_datafile
+
+    def _fix_old_file_paths(self, path):
+        if type(path) == float:  # pd.read_csv turns blank into np.nan
+            return path
+        if re.match(_EDGAR_RE, path) is None:
+            path = _EDGAR + path
+        return path
+
+    def read(self):
+        try:
+            self._sec_ftp_session = FTP(_SEC_FTP)
+            self._sec_ftp_session.login()
+        except EOFError:
+            raise RemoteDataError('FTP server has closed the connection.')
+        try:
+            if self.symbols == 'full':
+                return self._read_one_data(self.url, self.params)
+
+            elif self.symbols == 'daily':
+                return self._read_daily_data(self.url, self.params)
+        finally:
+            self._sec_ftp_session.quit()
+
+    def _sanitize_dates(self, start, end):
+        if is_number(start):
+            start = dt.datetime(start, 1, 1)
+        start = to_datetime(start)
+
+        if is_number(end):
+            end = dt.datetime(end, 1, 1)
+        end = to_datetime(end)
+
+        if start is None:
+            start = dt.datetime(2015, 1, 1)
+        if end is None:
+            end = dt.datetime(2015, 1, 3)
+        if start < _EDGAR_MIN_DATE:
+            start = _EDGAR_MIN_DATE
+
+        return start, end
+
+    def _get_dir_lists(self):
+        mlsd_tree = self._get_mlsd_tree(_EDGAR_DAILY)
+        return mlsd_tree
+
+    def _get_mlsd_tree(self, dir, top=True):
+        initial_mlsd = self._get_mlsd(dir)
+        mlsd = initial_mlsd[:]
+        for entry in initial_mlsd:
+            if entry['type'] == 'dir':
+                if top is True:
+                    if self._check_mlsd_year(entry) is not True:
+                        continue
+                subdir = dir + '/' + entry['name']
+                mlsd.extend(self._get_mlsd_tree(subdir, False))
+        return mlsd
+
+    def _get_mlsd(self, dir):
+        dir_list = []
+        try:
+            self._sec_ftp_session.retrlines('MLSD' + ' ' + dir,
+                                            dir_list.append)
+        except EOFError:
+            raise RemoteDataError('FTP server has closed the connection.')
+
+        dict_list = []
+        for line in dir_list:
+            entry = self._process_mlsd_line(line)
+            entry['path'] = dir
+            dict_list.append(entry)
+
+        return dict_list
+
+    def _process_mlsd_line(self, line):
+        line_dict = re.match(_MLSD_VALUES_RE, line).groupdict()
+        line_dict['date'] = self._get_index_date(line_dict['name'])
+        return line_dict
+
+    def _get_index_date(self, filename):
+        try:
+            idx_date = re.search(_FILENAME_DATE_RE, filename).group(1)
+            if len(idx_date) == 6:
+                if idx_date[-2:] == '94':
+                    filedate = dt.datetime.strptime(idx_date, '%m%d%y')
+                else:
+                    filedate = dt.datetime.strptime(idx_date, '%y%m%d')
+                    if filedate > _EDGAR_MAX_6_DIGIT_DATE:
+                        filedate = None
+            elif len(idx_date) == 8:
+                filedate = dt.datetime.strptime(idx_date, '%Y%m%d')
+        except AttributeError:
+            filedate = None
+
+        return filedate
+
+    def _check_mlsd_year(self, entry):
+        try:
+            if (self.start.year <= int(entry['name']) <= self.end.year):
+                return True
+            else:
+                return False
+        except TypeError:
+            return False
diff --git a/pandas_datareader/tests/test_data.py b/pandas_datareader/tests/test_data.py
@@ -528,10 +528,6 @@ def test_read_fred(self):
         vix = DataReader("VIXCLS", "fred")
         assert isinstance(vix, DataFrame)
 
-    def test_read_edgar_index(self):
-        ed = DataReader("full", "edgar-index")
-        assert isinstance(ed, DataFrame)
-
     def test_not_implemented(self):
         self.assertRaises(NotImplementedError, DataReader, "NA", "NA")
 
diff --git a/pandas_datareader/tests/test_edgar.py b/pandas_datareader/tests/test_edgar.py
@@ -2,13 +2,40 @@
 import pandas.util.testing as tm
 
 import pandas_datareader.data as web
+from pandas_datareader._utils import RemoteDataError
 
 
 class TestEdgarIndex(tm.TestCase):
-    def test_get_index(self):
-        ed = web.DataReader('full', 'edgar-index')
+    def test_get_full_index(self):
+        try:
+            ed = web.DataReader('full', 'edgar-index')
+        except RemoteDataError as e:
+            raise nose.SkipTest(e)
         assert len(ed > 1000)
 
+    def test_get_nonzip_index_and_low_date(self):
+        try:
+            ed = web.DataReader('daily', 'edgar-index', '1994-06-30',
+                                '1994-07-02')
+        except RemoteDataError as e:
+            raise nose.SkipTest(e)
+        assert len(ed > 200)
+
+    def test_get_gz_index_and_no_date(self):
+        try:
+            ed = web.DataReader('daily', 'edgar-index')
+        except RemoteDataError as e:
+            raise nose.SkipTest(e)
+        assert len(ed > 2000)
+
+    def test_6_digit_date(self):
+        try:
+            ed = web.DataReader('daily', 'edgar-index', '1998-05-18',
+                                '1998-05-18')
+        except RemoteDataError as e:
+            raise nose.SkipTest(e)
+        assert len(ed < 1200)
+
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                    exit=False)
