[ENH] Fama/French

* add get_datasets_famafrench (requires bs4)
* complete rewrite of get_data_famafrench

Author: 0x0L

pandas_datareader/data.py

@@ -8,6 +8,7 @@
 import datetime as dt
 import time
 import csv
+import re
 
 from collections import defaultdict
 
@@ -581,41 +582,117 @@ def fetch_data(url, name):
     return df
 
 
-_FAMAFRENCH_URL = 'http://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp'
+_FAMAFRENCH_URL = 'http://mba.tuck.dartmouth.edu/pages/faculty/ken.french/'
+_FF_PREFIX = 'ftp/'
+_FF_SUFFIX = '_CSV.zip'
 
 
-def get_data_famafrench(name):
-    # path of zip files
-    zip_file_path = '{0}/{1}_TXT.zip'.format(_FAMAFRENCH_URL, name)
+def get_datasets_famafrench():
+    """
+    Get the list of datasets available from the Fama/French data library.
+
+    Returns
+    -------
+    A list of valid inputs for get_data_famafrench.
+    """
+    from bs4 import BeautifulSoup
+
+    with urlopen(_FAMAFRENCH_URL + 'data_library.html') as socket:
+        root = BeautifulSoup(socket.read(), 'html.parser')
+
+    l = filter(lambda x: x.startswith(_FF_PREFIX) and x.endswith(_FF_SUFFIX),
+               [e.attrs['href'] for e in root.findAll('a') if 'href' in e.attrs])
 
-    with urlopen(zip_file_path) as url:
-        raw = url.read()
+    return list(map(lambda x: x[len(_FF_PREFIX):-len(_FF_SUFFIX)], l))
+
+
+def _download_data_famafrench(name):
+    url = '{0}{1}{2}{3}'.format(_FAMAFRENCH_URL, _FF_PREFIX, name, _FF_SUFFIX)
+    with urlopen(url) as socket:
+        raw = socket.read()
 
     with tempfile.TemporaryFile() as tmpf:
         tmpf.write(raw)
 
         with ZipFile(tmpf, 'r') as zf:
-            data = zf.open(zf.namelist()[0]).readlines()
-
-    line_lengths = np.array(lmap(len, data))
-    file_edges = np.where(line_lengths == 2)[0]
-
-    datasets = {}
-    edges = zip(file_edges + 1, file_edges[1:])
-    for i, (left_edge, right_edge) in enumerate(edges):
-        dataset = [d.split() for d in data[left_edge:right_edge]]
-        if len(dataset) > 10:
-            ncol_raw = np.array(lmap(len, dataset))
-            ncol = np.median(ncol_raw)
-            header_index = np.where(ncol_raw == ncol - 1)[0][-1]
-            header = dataset[header_index]
-            ds_header = dataset[header_index + 1:]
-            # to ensure the header is unique
-            header = ['{0} {1}'.format(j, hj) for j, hj in enumerate(header,
-                                                                     start=1)]
-            index = np.array([d[0] for d in ds_header], dtype=int)
-            dataset = np.array([d[1:] for d in ds_header], dtype=float)
-            datasets[i] = DataFrame(dataset, index, columns=header)
+            data = zf.open(zf.namelist()[0]).read().decode()
+
+    return data
+
+
+def _parse_date_famafrench(x):
+    # what's the correct python way to do that ??
+    x = x.strip()
+    try: return dt.datetime.strptime(x, '%Y')
+    except: pass
+    try: return dt.datetime.strptime(x, '%Y%m')
+    except: pass
+    return to_datetime(x)
+
+
+def get_data_famafrench(name):
+    """
+    Get data for the given name from the Fama/French data library.
+
+    For annual and monthly data, index is a pandas.PeriodIndex, otherwise
+    it's a pandas.DatetimeIndex.
+
+    Returns
+    -------
+    df : a dictionary of pandas.DataFrame. Tables are accessed by integer keys.
+         See df['DESCR'] for a description of the dataset
+    """
+    params = {'index_col': 0,
+              'parse_dates': [0],
+              'date_parser': _parse_date_famafrench}
+
+    if name.endswith('_Breakpoints'):
+        # headers in these files are not valid
+        params['skiprows'] = 1
+        if name == 'Prior_2-12_Breakpoints':
+            params['skiprows'] = 3
+
+        params['names'] = ['Date']
+        if name.find('-') > -1:
+            params['names'] += ['<=0', '>0']
+        else:
+            params['names'] += ['Count']
+        r = range(0, 105, 5)
+        params['names'] += list(zip(r, r[1:]))
+
+    data = _download_data_famafrench(name)
+
+    # chunks are split by blank lines
+    # chunks with more than 800 characters are considered data
+    # those with less characters are considered documentation
+    doc, almost_csv = [], []
+    for chunk in data.split(2 * '\r\n'):
+        if len(chunk) < 800:
+            doc.append(chunk.replace('\r\n', ' ').strip())
+        else:
+            almost_csv.append(chunk)
+    desc = ' '.join(doc).replace(2*' ', ' ')
+
+    datasets, table_desc = {}, []
+    for i, src in enumerate(almost_csv):
+        # the table starts there
+        match = re.search('^\s*,', src, re.M)
+        start = 0 if not match else match.start()
+
+        df = read_csv(StringIO('Date' + src[start:]), **params)
+        try: df = df.to_period()
+        except: pass
+        datasets[i] = df
+
+        short_desc = ' '.join([src[:start].replace('\r\n', ' ').strip(),
+                               '({0} rows x {1} cols)'.format(*df.shape)])
+        table_desc.append(short_desc)
+
+    table_str = map(lambda x: '{} : {}'.format(*x), enumerate(table_desc))
+
+    TEMPLATE = '{0}\n{1}\n\n{2}\n\n{3}'
+    datasets['DESCR'] = TEMPLATE.format(name.replace('_', ' '), len(name)*'-',
+                                        desc, '\n'.join(table_str))
 
     return datasets