Merge pull request #501 from sgsmob/export

Expand functionality of `delphi_utils.create_export_csv()` for use in more indicators

Author: krivard

_delphi_utils_python/delphi_utils/export.py

@@ -2,16 +2,19 @@
 # -*- coding: utf-8 -*-
 from datetime import datetime
 from os.path import join
+from typing import Optional
 
+import numpy as np
 import pandas as pd
 
 def create_export_csv(
     df: pd.DataFrame,
-    start_date: datetime,
     export_dir: str,
-    metric: str,
     geo_res: str,
     sensor: str,
+    metric: Optional[str] = None,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None
 ):
     """Export data in the format expected by the Delphi API.
 
@@ -21,21 +24,36 @@ def create_export_csv(
         Columns: geo_id, timestamp, val, se, sample_size
     export_dir: str
         Export directory
-    metric: str
-        Metric we are considering
     geo_res: str
         Geographic resolution to which the data has been aggregated
     sensor: str
         Sensor that has been calculated (cumulative_counts vs new_counts)
+    metric: Optional[str]
+        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:
-        export_fn = f'{date.strftime("%Y%m%d")}_{geo_res}_' f"{metric}_{sensor}.csv"
-        export_file = join(export_dir, export_fn)
+        if metric is None:
+            export_filename = f"{date.strftime('%Y%m%d')}_{geo_res}_{sensor}.csv"
+        else:
+            export_filename = f"{date.strftime('%Y%m%d')}_{geo_res}_{metric}_{sensor}.csv"
+        export_file = join(export_dir, export_filename)
         df[df["timestamp"] == date][["geo_id", "val", "se", "sample_size",]].to_csv(
             export_file, index=False, na_rep="NA"
         )

_delphi_utils_python/tests/test_dir/.gitignore

@@ -0,0 +1 @@
+*.csv

_delphi_utils_python/tests/test_dir/20200215_county_deaths_test.csv

@@ -1,47 +1,155 @@
-import pytest
-
+"""Tests for exporting CSV files."""
 from datetime import datetime
 from os import listdir, remove
 from os.path import join
 
 import pandas as pd
 from delphi_utils import create_export_csv
 
+def _clean_directory(directory):
+    """Clean files out of a directory."""
+    for fname in listdir(directory):
+        if fname.startswith("."):
+            continue
+        remove(join(directory, fname))
+
+
+def _non_ignored_files_set(directory):
+    """List all files in a directory not preceded by a '.' and store them in a set."""
+    out = set()
+    for fname in listdir(directory):
+        if fname.startswith("."):
+            continue
+        out.add(fname)
+    return out
+
 
 class TestExport:
-    def test_export_csv(self):
+    """Tests for exporting CSVs."""
+    # List of times for data points.
+    TIMES = [
+        datetime.strptime(x, "%Y-%m-%d")
+        for x in ["2020-02-15", "2020-02-15", "2020-03-01", "2020-03-15"]
+    ]
+
+    # A sample data frame.
+    DF = pd.DataFrame(
+        {
+            "geo_id": ["51093", "51175", "51175", "51620"],
+            "timestamp": TIMES,
+            "val": [3.6, 2.1, 2.2, 2.6],
+            "se": [0.15, 0.22, 0.20, 0.34],
+            "sample_size": [100, 100, 101, 100],
+        }
+    )
+
+    # Directory in which to store tests.
+    TEST_DIR = "test_dir"
+
+    def test_export_with_metric(self):
+        """Test that exporting CSVs with the `metrics` argument yields the correct files."""
 
         # Clean receiving directory
-        for fname in listdir("test_dir"):
-            remove(join("test_dir", fname))
-
-        times = [
-            datetime.strptime(x, "%Y-%m-%d")
-            for x in ["2020-02-15", "2020-02-15", "2020-03-01", "2020-03-15"]
-        ]
-        df = pd.DataFrame(
-            {
-                "geo_id": ["51093", "51175", "51175", "51620"],
-                "timestamp": times,
-                "val": [3.6, 2.1, 2.2, 2.6],
-                "se": [0.15, 0.22, 0.20, 0.34],
-                "sample_size": [100, 100, 101, 100],
-            }
-        )
+        _clean_directory(self.TEST_DIR)
 
         create_export_csv(
-            df=df,
+            df=self.DF,
             start_date=datetime.strptime("2020-02-15", "%Y-%m-%d"),
-            export_dir="test_dir",
+            export_dir=self.TEST_DIR,
             metric="deaths",
             geo_res="county",
             sensor="test",
         )
 
-        assert set(listdir("test_dir")) == set(
+        assert _non_ignored_files_set(self.TEST_DIR) == set(
             [
                 "20200215_county_deaths_test.csv",
                 "20200301_county_deaths_test.csv",
                 "20200315_county_deaths_test.csv",
             ]
         )
+
+    def test_export_without_metric(self):
+        """Test that exporting CSVs without the `metrics` argument yields the correct files."""
+
+        # Clean receiving directory
+        _clean_directory(self.TEST_DIR)
+
+        create_export_csv(
+            df=self.DF,
+            start_date=datetime.strptime("2020-02-15", "%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",
+                "20200315_county_test.csv",
+            ]
+        )
+
+    def test_export_with_limiting_start_date(self):
+        """Test that the `start_date` prevents earlier dates from being exported."""
+
+        # Clean receiving directory
+        _clean_directory(self.TEST_DIR)
+
+        create_export_csv(
+            df=self.DF,
+            start_date=datetime.strptime("2020-02-20", "%Y-%m-%d"),
+            export_dir=self.TEST_DIR,
+            geo_res="county",
+            sensor="test",
+        )
+
+        assert _non_ignored_files_set(self.TEST_DIR) == set(
+            [
+                "20200301_county_test.csv",
+                "20200315_county_test.csv",
+            ]
+        )
+
+    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)
+
+        create_export_csv(
+            df=self.DF,
+            export_dir=self.TEST_DIR,
+            geo_res="state",
+            sensor="test",
+        )
+
+        assert _non_ignored_files_set(self.TEST_DIR) == set(
+            [
+                "20200215_state_test.csv",
+                "20200301_state_test.csv",
+                "20200315_state_test.csv",
+            ]
+        )