Merge pull request #449 from bashtage/kurtforrester-base-freq

Kurtforrester base freq

Author: bashtage

pandas_datareader/base.py

@@ -19,25 +19,29 @@ class _BaseReader(object):
     """
     Parameters
     ----------
-        sym : string with a single Single stock symbol (ticker).
+        symbols : {str, List[str]}
+            String symbol of like of symbols
         start : string, (defaults to '1/1/2010')
-                Starting date, timestamp. Parses many different kind of date
-                representations (e.g., 'JAN-01-2010', '1/1/10', 'Jan, 1, 1980')
+            Starting date, timestamp. Parses many different kind of date
+            representations (e.g., 'JAN-01-2010', '1/1/10', 'Jan, 1, 1980')
         end : string, (defaults to today)
-                Ending date, timestamp. Same format as starting date.
+            Ending date, timestamp. Same format as starting date.
         retry_count : int, default 3
-                Number of times to retry query request.
+            Number of times to retry query request.
         pause : float, default 0.1
-                Time, in seconds, of the pause between retries.
+            Time, in seconds, of the pause between retries.
         session : Session, default None
-                requests.sessions.Session instance to be used
+            requests.sessions.Session instance to be used
+        freq : {str, None}
+            Frequency to use in select readers
     """
 
     _chunk_size = 1024 * 1024
     _format = 'string'
 
-    def __init__(self, symbols, start=None, end=None,
-                 retry_count=3, pause=0.1, timeout=30, session=None):
+    def __init__(self, symbols, start=None, end=None,  retry_count=3,
+                 pause=0.1, timeout=30, session=None, freq=None):
+
         self.symbols = symbols
 
         start, end = _sanitize_dates(start, end)
@@ -51,6 +55,7 @@ def __init__(self, symbols, start=None, end=None,
         self.timeout = timeout
         self.pause_multiplier = 1
         self.session = _init_session(session, retry_count)
+        self.freq = freq
 
     def close(self):
         """ close my session """

pandas_datareader/tests/test_wb.py

@@ -1,14 +1,14 @@
 import time
-import pytest
 
 import numpy as np
 import pandas as pd
+import pandas.util.testing as tm
+import pytest
 import requests
 
-import pandas.util.testing as tm
+from pandas_datareader.compat import assert_raises_regex
 from pandas_datareader.wb import (search, download, get_countries,
                                   get_indicators, WorldBankReader)
-from pandas_datareader.compat import assert_raises_regex
 
 
 class TestWB(object):
@@ -214,3 +214,64 @@ def test_wdi_get_indicators(self):
             # assert_index_equal doesn't exists
             assert result.columns.equals(exp_col)
             assert len(result) > 10000
+
+    def test_wdi_download_monthly(self):
+        expected = {'COPPER': {('World', '2012M01'): 8040.47,
+                               ('World', '2011M12'): 7565.48,
+                               ('World', '2011M11'): 7581.02,
+                               ('World', '2011M10'): 7394.19,
+                               ('World', '2011M09'): 8300.14,
+                               ('World', '2011M08'): 9000.76,
+                               ('World', '2011M07'): 9650.46,
+                               ('World', '2011M06'): 9066.85,
+                               ('World', '2011M05'): 8959.90,
+                               ('World', '2011M04'): 9492.79,
+                               ('World', '2011M03'): 9503.36,
+                               ('World', '2011M02'): 9867.60,
+                               ('World', '2011M01'): 9555.70}}
+        expected = pd.DataFrame(expected)
+        # Round, to ignore revisions to data.
+        expected = np.round(expected, decimals=-3)
+        expected = expected.sort_index()
+        cntry_codes = 'ALL'
+        inds = 'COPPER'
+        result = download(country=cntry_codes, indicator=inds,
+                          start=2011, end=2012, freq='M', errors='ignore')
+        result = result.sort_index()
+        result = np.round(result, decimals=-3)
+
+        expected.index.names = ['country', 'year']
+        tm.assert_frame_equal(result, expected)
+
+        result = WorldBankReader(inds, countries=cntry_codes, start=2011,
+                                 end=2012, freq='M', errors='ignore').read()
+        result = result.sort_index()
+        result = np.round(result, decimals=-3)
+        tm.assert_frame_equal(result, expected)
+
+    def test_wdi_download_quarterly(self):
+        code = 'DT.DOD.PUBS.CD.US'
+        expected = {code: {('Albania', '2012Q1'): 3240539817.18,
+                           ('Albania', '2011Q4'): 3213979715.15,
+                           ('Albania', '2011Q3'): 3187681048.95,
+                           ('Albania', '2011Q2'): 3248041513.86,
+                           ('Albania', '2011Q1'): 3137210567.92}}
+        expected = pd.DataFrame(expected)
+        # Round, to ignore revisions to data.
+        expected = np.round(expected, decimals=-3)
+        expected = expected.sort_index()
+        cntry_codes = 'ALB'
+        inds = 'DT.DOD.PUBS.CD.US'
+        result = download(country=cntry_codes, indicator=inds,
+                          start=2011, end=2012, freq='Q', errors='ignore')
+        result = result.sort_index()
+        result = np.round(result, decimals=-3)
+
+        expected.index.names = ['country', 'year']
+        tm.assert_frame_equal(result, expected)
+
+        result = WorldBankReader(inds, countries=cntry_codes, start=2011,
+                                 end=2012, freq='Q', errors='ignore').read()
+        result = result.sort_index()
+        result = np.round(result, decimals=-1)
+        tm.assert_frame_equal(result, expected)

pandas_datareader/wb.py

@@ -116,7 +116,7 @@ class WorldBankReader(_BaseReader):
     _format = 'json'
 
     def __init__(self, symbols=None, countries=None,
-                 start=None, end=None,
+                 start=None, end=None, freq=None,
                  retry_count=3, pause=0.001, session=None, errors='warn'):
 
         if symbols is None:
@@ -144,6 +144,14 @@ def __init__(self, symbols=None, countries=None,
                 warnings.warn('Non-standard ISO '
                               'country codes: %s' % tmp, UserWarning)
 
+        freq_symbols = ['M', 'Q', 'A', None]
+
+        if freq not in freq_symbols:
+            msg = 'The frequency `{0}` is not in the accepted ' \
+                  'list.'.format(freq)
+            raise ValueError(msg)
+
+        self.freq = freq
         self.countries = countries
         self.errors = errors
 
@@ -155,8 +163,18 @@ def url(self):
 
     @property
     def params(self):
-        return {'date': '{0}:{1}'.format(self.start.year, self.end.year),
-                'per_page': 25000, 'format': 'json'}
+        if self.freq == 'M':
+            return {'date': '{0}M{1:02d}:{2}M{3:02d}'.format(self.start.year,
+                    self.start.month, self.end.year, self.end.month),
+                    'per_page': 25000, 'format': 'json'}
+        elif self.freq == 'Q':
+            return {'date': '{0}Q{1}:{2}Q{3}'.format(self.start.year,
+                    self.start.quarter, self.end.year,
+                    self.end.quarter), 'per_page': 25000,
+                    'format': 'json'}
+        else:
+            return {'date': '{0}:{1}'.format(self.start.year, self.end.year),
+                    'per_page': 25000, 'format': 'json'}
 
     def read(self):
         try:
@@ -331,7 +349,7 @@ def search(self, string='gdp.*capi', field='name', case=False):
         return out
 
 
-def download(country=None, indicator=None, start=2003, end=2005,
+def download(country=None, indicator=None, start=2003, end=2005, freq=None,
              errors='warn', **kwargs):
     """
     Download data series from the World Bank's World Development Indicators
@@ -357,6 +375,11 @@ def download(country=None, indicator=None, start=2003, end=2005,
     end: int
         Last year of the data series (inclusive)
 
+    freq: str
+        frequency or periodicity of the data to be retrieved (e.g. 'M' for
+        monthly, 'Q' for quarterly, and 'A' for annual). None defaults to
+        annual.
+
     errors: str {'ignore', 'warn', 'raise'}, default 'warn'
         Country codes are validated against a hardcoded list.  This controls
         the outcome of that validation, and attempts to also apply
@@ -375,7 +398,7 @@ def download(country=None, indicator=None, start=2003, end=2005,
 
     """
     return WorldBankReader(symbols=indicator, countries=country,
-                           start=start, end=end, errors=errors,
+                           start=start, end=end, freq=freq, errors=errors,
                            **kwargs).read()