From eda64d1533bb6cc388fff701d2ba439c9e47c1a6 Mon Sep 17 00:00:00 2001
From: Samuel Gratzl <sam@sgratzl.com>
Date: Thu, 6 May 2021 13:42:28 -0400
Subject: [PATCH] feat: exclude fields

by using fields=-fieldname
---
 integrations/server/test_covidcast.py | 25 ++++++++++++++++++++++
 src/server/_query.py                  | 30 +++++++++++++++++++++------
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/integrations/server/test_covidcast.py b/integrations/server/test_covidcast.py
index 8aaa063d3..d5029fe4e 100644
--- a/integrations/server/test_covidcast.py
+++ b/integrations/server/test_covidcast.py
@@ -274,6 +274,31 @@ def test_fields(self):
       'message': 'success',
     })
 
+
+    # limit exclude fields
+    response = requests.get(BASE_URL, params={
+      'endpoint': 'covidcast',
+      'data_source': 'src',
+      'signal': 'sig',
+      'time_type': 'day',
+      'geo_type': 'county',
+      'time_values': 20200414,
+      'geo_value': '01234',
+      'fields': '-value,-stderr,-sample_size,-direction,-issue,-lag,-signal'
+    })
+    response.raise_for_status()
+    response = response.json()
+
+    # assert that the right data came back
+    self.assertEqual(response, {
+      'result': 1,
+      'epidata': [{
+        'time_value': 20200414,
+        'geo_value': '01234'
+       }],
+      'message': 'success',
+    })
+
   def test_location_wildcard(self):
     """Select all locations with a wildcard query."""
 
diff --git a/src/server/_query.py b/src/server/_query.py
index 2c595a4f0..1a915b85e 100644
--- a/src/server/_query.py
+++ b/src/server/_query.py
@@ -97,11 +97,21 @@ def filter_fields(generator: Iterable[Dict[str, Any]]):
     if not fields:
         yield from generator
     else:
+        exclude_fields = {f[1:] for f in fields if f.startswith("-")}
+        include_fields = [f for f in fields if not f.startswith("-") and f not in exclude_fields]
+
         for row in generator:
             filtered = dict()
-            for field in fields:
-                if field in row:
-                    filtered[field] = row[field]
+            if include_fields:
+                # positive list
+                for field in include_fields:
+                    if field in row:
+                        filtered[field] = row[field]
+            elif exclude_fields:
+                # negative list
+                for k, v in row.items():
+                    if k not in exclude_fields:
+                        filtered[k] = v
             yield filtered
 
 
@@ -252,9 +262,17 @@ def execute_queries(
 
     fields_to_send = set(extract_strings("fields") or [])
     if fields_to_send:
-        fields_string = [v for v in fields_string if v in fields_to_send]
-        fields_int = [v for v in fields_int if v in fields_to_send]
-        fields_float = [v for v in fields_float if v in fields_to_send]
+        exclude_fields = {f[1:] for f in fields_to_send if f.startswith("-")}
+        include_fields = {f for f in fields_to_send if not f.startswith("-") and f not in exclude_fields}
+
+        if include_fields:
+            fields_string = [v for v in fields_string if v in include_fields]
+            fields_int = [v for v in fields_int if v in include_fields]
+            fields_float = [v for v in fields_float if v in include_fields]
+        if exclude_fields:
+            fields_string = [v for v in fields_string if v not in exclude_fields]
+            fields_int = [v for v in fields_int if v not in exclude_fields]
+            fields_float = [v for v in fields_float if v not in exclude_fields]
 
     query_list = list(queries)