add end date functionality to export csvs

Author: sgsmob

_delphi_utils_python/delphi_utils/export.py

@@ -4,15 +4,17 @@
 from os.path import join
 from typing import Optional
 
+import numpy as np
 import pandas as pd
 
 def create_export_csv(
     df: pd.DataFrame,
     export_dir: str,
     geo_res: str,
     sensor: str,
+    metric: Optional[str] = None,
     start_date: Optional[datetime] = None,
-    metric: Optional[str] = None
+    end_date: Optional[datetime] = None
 ):
     """Export data in the format expected by the Delphi API.
 
@@ -26,18 +28,24 @@ def create_export_csv(
         Geographic resolution to which the data has been aggregated
     sensor: str
         Sensor that has been calculated (cumulative_counts vs new_counts)
-    start_date: Optional[datetime]
-        Earliest date to export.  If None, all dates are exported.
     metric: Optional[str]
-        Metric we are considering, if any
+        Metric we are considering, if any.
+    start_date: Optional[datetime]
+        Earliest date to export or None if no minimum date restrictions should be applied.
+    end_date: Optional[datetime]
+        Latest date to export or None if no maximum date restrictions should be applied.
     """
     df = df.copy()
 
     df["timestamp"] = pd.to_datetime(df["timestamp"])
     if start_date is None:
         start_date = min(df["timestamp"])
+    if end_date is None:
+        end_date = max(df["timestamp"])
+
     dates = pd.Series(
-        df[df["timestamp"] >= start_date]["timestamp"].unique()
+        df[np.logical_and(df["timestamp"] >= start_date,
+                          df["timestamp"] <= end_date)]["timestamp"].unique()
     ).sort_values()
 
     for date in dates:

_delphi_utils_python/tests/test_export.py

@@ -112,8 +112,29 @@ def test_export_with_limiting_start_date(self):
             ]
         )
 
-    def test_export_with_no_start_date(self):
-        """Test that omitting the `start_date` exports all dates."""
+    def test_export_with_limiting_end_date(self):
+        """Test that the `end_date` prevents later dates from being exported."""
+
+        # Clean receiving directory
+        _clean_directory(self.TEST_DIR)
+
+        create_export_csv(
+            df=self.DF,
+            end_date=datetime.strptime("2020-03-07", "%Y-%m-%d"),
+            export_dir=self.TEST_DIR,
+            geo_res="county",
+            sensor="test",
+        )
+
+        assert _non_ignored_files_set(self.TEST_DIR) == set(
+            [
+                "20200215_county_test.csv",
+                "20200301_county_test.csv",
+            ]
+        )
+
+    def test_export_with_no_dates(self):
+        """Test that omitting the `start_date` and `end_date` exports all dates."""
 
         # Clean receiving directory
         _clean_directory(self.TEST_DIR)