ENH: Adds support for Enigma datasets

Init commit

Cleans up code

Improves failure messaging

Adds usage, improves errors

Adds docs

Bumps version

Adds retry logic, improves docs

Fixes import

Adds DataReader as an entry point for enigma

Adds query caching

Formatting

Reads DataReader before returning

Adds tests

Removes unused import

Cleans up code

Fixes test

Cleans up unused imports

Updates User-Agent header

Removes hardcoded API token

Author: Trevor Prater

docs/source/remote_data.rst

@@ -25,6 +25,7 @@ Currently the following sources are supported:
 
     - :ref:`Yahoo! Finance<remote_data.yahoo>`
     - :ref:`Google Finance<remote_data.google>`
+    - :ref:`Enigma<remote_data.enigma>`
     - :ref:`St.Louis FED (FRED)<remote_data.fred>`
     - :ref:`Kenneth French's data library<remote_data.ff>`
     - :ref:`World Bank<remote_data.wb>`
@@ -195,6 +196,22 @@ Available expiry dates can be accessed from the ``expiry_dates`` property.
       data = goog.get_options_data(expiry=goog.expiry_dates[0])
       data.iloc[0:5, 0:5]
 
+.. _remote_data.enigma:
+
+Enigma
+======
+
+Access datasets from `Enigma <https://app.enigma.io>`__,
+the world's largest repository of structured public data.
+
+.. ipython:: python
+
+    import os
+    import pandas_datareader as pdr
+
+    df = pdr.get_data_enigma('enigma.trade.ams.toxic.2015', os.getenv('ENIGMA_API_KEY'))
+    df.columns
+
 .. _remote_data.fred:
 
 FRED

pandas_datareader/__init__.py

@@ -1,4 +1,4 @@
-__version__ = version = '0.2.1'
+__version__ = version = '0.2.2'
 
-from .data import (get_components_yahoo, get_data_famafrench, get_data_google, get_data_yahoo,  # noqa
+from .data import (get_components_yahoo, get_data_famafrench, get_data_google, get_data_yahoo, get_data_enigma,  # noqa
         get_data_yahoo_actions, get_quote_google, get_quote_yahoo, DataReader, Options)  # noqa

pandas_datareader/data.py

@@ -19,6 +19,7 @@
 from pandas_datareader.famafrench import FamaFrenchReader
 from pandas_datareader.oecd import OECDReader
 from pandas_datareader.edgar import EdgarIndexReader
+from pandas_datareader.enigma import EnigmaReader
 from pandas_datareader.oanda import get_oanda_currency_historical_rates
 
 
@@ -38,6 +39,10 @@ def get_data_yahoo(*args, **kwargs):
     return YahooDailyReader(*args, **kwargs).read()
 
 
+def get_data_enigma(*args, **kwargs):
+    return EnigmaReader(*args, **kwargs).read()
+
+
 def get_data_yahoo_actions(*args, **kwargs):
     return YahooActionReader(*args, **kwargs).read()
 
@@ -51,7 +56,7 @@ def get_quote_google(*args, **kwargs):
 
 
 def DataReader(name, data_source=None, start=None, end=None,
-               retry_count=3, pause=0.001, session=None):
+               retry_count=3, pause=0.001, session=None, access_key=None):
     """
     Imports data from a number of online sources.
 
@@ -125,6 +130,9 @@ def DataReader(name, data_source=None, start=None, end=None,
                                  retry_count=retry_count, pause=pause,
                                  session=session).read()
 
+    elif data_source == "enigma":
+        return EnigmaReader(datapath=name, api_key=access_key).read()
+
     elif data_source == "fred":
         return FredReader(symbols=name, start=start, end=end,
                           retry_count=retry_count, pause=pause,

pandas_datareader/enigma.py

@@ -0,0 +1,107 @@
+import zlib
+import json
+import os
+import sys
+import time
+from pandas.compat import StringIO
+
+from pandas import DataFrame
+import pandas.compat as compat
+import pandas as pd
+import requests
+
+from pandas_datareader.base import _BaseReader
+
+
+class EnigmaReader(_BaseReader):
+    """
+    Collects Enigma data located at the specified datapath and returns a pandas DataFrame.
+
+    Usage (high-level):
+    ```
+        import pandas_datareader as pdr
+        df = pdr.get_data_enigma('enigma.inspections.restaurants.fl')
+
+        #in the event that ENIGMA_API_KEY does not exist in your env, it can be supplied as the second arg:
+        df = prd.get_data_enigma('enigma.inspections.restaurants.fl', 'ARIAMFHKJMISF38UT')
+    ```
+
+    Usage:
+    ```
+        df = EnigmaReader(datapath='enigma.inspections.restaurants.fl', api_key='ARIAMFHKJMISF38UT').read()
+    ```
+    """
+
+    def __init__(self,
+                 datapath=None,
+                 api_key=None,
+                 retry_count=5,
+                 pause=0.250,
+                 session=None):
+
+        super(EnigmaReader, self).__init__(symbols=[],
+                                           retry_count=retry_count,
+                                           pause=pause)
+        if api_key == None:
+            self._api_key = os.getenv('ENIGMA_API_KEY')
+            if self._api_key == None:
+                raise ValueError(
+                    """Please provide an Enigma API key or set the ENIGMA_API_KEY environment variable\n
+                        If you do not have an API key, you can get one here: https://app.enigma.io/signup""")
+        else:
+            self._api_key = api_key
+
+        self._datapath = datapath
+        if not isinstance(self._datapath, compat.string_types):
+            raise ValueError(
+                "The Enigma datapath must be a string (ex: 'enigma.inspections.restaurants.fl')")
+
+            
+    @property
+    def url(self):
+        return 'https://api.enigma.io/v2/export/{}/{}'.format(self._api_key,
+                                                              self._datapath)
+
+    @property
+    def export_key(self):
+        return 'export_url'
+
+    
+    @property
+    def _head_key(self):
+        return 'head_url'
+
+    
+    def _request(self, url):
+        self.session.headers.update({'User-Agent': 'pandas-datareader'})
+        resp = self.session.get(url)
+        resp.raise_for_status()
+        return resp
+
+    
+    def _decompress_export(self, compressed_export_data):
+        return zlib.decompress(compressed_export_data, 16 + zlib.MAX_WBITS)
+
+    
+    def extract_export_url(self, delay=10, max_attempts=10):
+        """
+        Performs an HTTP HEAD request on 'head_url' until it returns a `200`.
+        This allows the Enigma API time to export the requested data.
+        """
+        resp = self._request(self.url)
+        attempts = 0
+        while True:
+            try:
+                requests.head(resp.json()[self._head_key]).raise_for_status()
+            except Exception as e:
+                attempts += 1
+                if attempts > max_attempts:
+                    raise e
+                time.sleep(delay)
+                continue
+            return resp.json()[self.export_key]
+
+    def read(self):
+        export_gzipped_req = self._request(self.extract_export_url())
+        decompressed_data = self._decompress_export(export_gzipped_req.content)
+        return pd.read_csv(StringIO(decompressed_data))

pandas_datareader/tests/test_enigma.py

@@ -0,0 +1,54 @@
+import os
+
+import requests
+from requests.exceptions import HTTPError
+
+import nose
+import pandas.util.testing as tm
+from pandas.util.testing import (assert_series_equal, assert_frame_equal)
+from pandas_datareader.tests._utils import _skip_if_no_lxml
+
+import pandas_datareader.data as web
+import pandas_datareader as pdr
+
+TEST_API_KEY = os.getenv('ENIGMA_API_KEY')
+
+
+class TestEnigma(tm.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        super(TestEnigma, cls).setUpClass()
+        _skip_if_no_lxml()
+
+    def test_enigma(self):
+        self.assertTrue('serialid' in list(
+            web.DataReader('enigma.inspections.restaurants.fl',
+                           'enigma',
+                           access_key=TEST_API_KEY).columns))
+        self.assertTrue('serialid' in list(pdr.get_data_enigma(
+            'enigma.inspections.restaurants.fl', TEST_API_KEY)))
+
+    def test_bad_key(self):
+        _exception = None
+        try:
+            web.DataReader('enigma.inspections.restaurants.fl',
+                           'enigma',
+                           access_key=TEST_API_KEY + 'xxx')
+        except HTTPError as e:
+            _exception = e
+        assert isinstance(_exception, HTTPError)
+
+    def test_bad_url(self):
+        _exception = None
+        try:
+            web.DataReader('enigma.inspections.restaurants.fllzzy',
+                           'enigma',
+                           access_key=TEST_API_KEY)
+        except Exception as e:
+            _exception = e
+        assert isinstance(_exception, HTTPError)
+
+
+if __name__ == '__main__':
+    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
+                   exit=False)