Merge pull request #396 from chinandrew/epidata-async

Add async method to epidata

Author: krivard

integrations/client/test_delphi_epidata.py

@@ -471,3 +471,48 @@ def test_covidcast_nowcast(self):
       'src', 'sig1', 'sensor', 'day', 'county', 22222222, '01001')
 
     self.assertEqual(response, {'result': -2, 'message': 'no results'})
+
+  def test_async_epidata(self):
+    # insert dummy data
+    self.cur.execute('''
+      insert into covidcast values
+        (0, 'src', 'sig', 'day', 'county', 20200414, '11111',
+          123, 10, 11, 12, 456, 13, 20200414, 0, 1, False),
+        (0, 'src', 'sig', 'day', 'county', 20200414, '22222',
+          123, 20, 21, 22, 456, 23, 20200414, 0, 1, False),
+        (0, 'src', 'sig', 'day', 'county', 20200414, '33333',
+          123, 30, 31, 32, 456, 33, 20200414, 0, 1, False),
+        (0, 'src', 'sig', 'day', 'msa', 20200414, '11111',
+          123, 40, 41, 42, 456, 43, 20200414, 0, 1, False),
+        (0, 'src', 'sig', 'day', 'msa', 20200414, '22222',
+          123, 50, 51, 52, 456, 53, 20200414, 0, 1, False),
+        (0, 'src', 'sig', 'day', 'msa', 20200414, '33333',
+          123, 60, 61, 62, 456, 634, 20200414, 0, 1, False)
+    ''')
+    self.cnx.commit()
+    test_output = Epidata.async_epidata([
+      {
+        'source': 'covidcast',
+        'data_source': 'src',
+        'signals': 'sig',
+        'time_type': 'day',
+        'geo_type': 'county',
+        'geo_value': '11111',
+        'time_values': '20200414'
+      },
+      {
+        'source': 'covidcast',
+        'data_source': 'src',
+        'signals': 'sig',
+        'time_type': 'day',
+        'geo_type': 'county',
+        'geo_value': '00000',
+        'time_values': '20200414'
+      }
+    ], batch_size=10)
+    responses = [i[0] for i in test_output]*12
+    # check response is same as standard covidcast call, using 24 calls to test batch sizing
+    self.assertEqual(responses,
+                     [Epidata.covidcast('src', 'sig', 'day', 'county', 20200414, '11111'),
+                      Epidata.covidcast('src', 'sig', 'day', 'county', 20200414, '00000')]*12
+                     )

src/client/delphi_epidata.py

@@ -10,7 +10,10 @@
 
 # External modules
 import requests
+import asyncio
+import warnings
 
+from aiohttp import ClientSession
 from pkg_resources import get_distribution, DistributionNotFound
 
 # Obtain package version for the user-agent. Uses the installed version by
@@ -704,3 +707,29 @@ def covidcast_nowcast(
 
     # Make the API call
     return Epidata._request(params)
+
+  @staticmethod
+  def async_epidata(param_list, batch_size=100):
+    """Make asynchronous Epidata calls for a list of parameters."""
+    async def async_get(params, session):
+      """Helper function to make Epidata GET requests."""
+      async with session.get(Epidata.BASE_URL, params=params) as response:
+        return await response.json(), params
+
+    async def async_make_calls(param_combos):
+      """Helper function to asynchronously make and aggregate Epidata GET requests."""
+      tasks = []
+      async with ClientSession() as session:
+        for param in param_combos:
+          task = asyncio.ensure_future(async_get(param, session))
+          tasks.append(task)
+        responses = await asyncio.gather(*tasks)
+        return responses
+
+    batches = [param_list[i:i+batch_size] for i in range(0, len(param_list), batch_size)]
+    responses = []
+    for batch in batches:
+      loop = asyncio.get_event_loop()
+      future = asyncio.ensure_future(async_make_calls(batch))
+      responses += loop.run_until_complete(future)
+    return responses

src/client/packaging/pypi/setup.py

@@ -14,6 +14,7 @@
   url='https://github.com/cmu-delphi/delphi-epidata',
   packages=setuptools.find_packages(),
   install_requires=[
+    'aiohttp'
     'requests>=2.7.0',
   ],
   classifiers=[