Add CSV/XLS export functionality

mkuthan · Nov 22, 2024 · ce24acd · ce24acd
1 parent dca3b29
commit ce24acd
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 20 deletions.
diff --git a/README.md b/README.md
@@ -35,6 +35,8 @@ This project demonstrates how to leverage the built-in power of Streamlit using
 * 🧪 Implement BigQuery integration tests
 * 📈 Add more visualizations for integrated public dataset
 * 🔐 Integration with external OAuth provider, see [roadmap](https://roadmap.streamlit.app/)
+* 📝 Add request logging
+* 📊 Add tabular data export to CSV
 * 🔄 Redirect to the original page after login
 * ⚖️ Describe load balancer strategies, for example: sticky session
 * 🔄 Automatically update dependencies, see [dependabot#10039](https://github.com/dependabot/dependabot-core/issues/10039)

diff --git a/example/services/nyc_tlc_trips.py b/example/services/nyc_tlc_trips.py
@@ -62,6 +62,11 @@ def get_trips(date_range: tuple[date, date], payment_type: str) -> pd.DataFrame:
     return results
 
 
+@st.cache_data(ttl=600)
+def export_trips(trips: pd.DataFrame) -> bytes:
+    return trips.to_csv(index=False).encode("utf-8")
+
+
 def _get_payment_type_key(payment_type: str) -> str:
     for key, value in _PAYMENT_TYPES.items():
         if value.lower() == payment_type.lower():

diff --git a/example/ui/pages/nyt_tlc_trips.py b/example/ui/pages/nyt_tlc_trips.py
@@ -16,7 +16,9 @@
 trips = nyc_tlc_trips.get_trips(date_range, payment_type)
 
 st.caption("Search results")
-st.dataframe(trips)
+st.dataframe(trips, use_container_width=True)
+
+st.download_button("Export CSV", nyc_tlc_trips.export_trips(trips), "trips.csv", "text/csv")
 
 trip_count_chart = (
     alt.Chart(trips)

diff --git a/tests/example/services/test_nyc_tlc_trips.py b/tests/example/services/test_nyc_tlc_trips.py
@@ -5,25 +5,26 @@
 
 from example.services import nyc_tlc_trips
 
+_any_trips = pd.DataFrame(
+    {
+        "day": ["2023-01-01", "2023-01-02"],
+        "payment_type": [1, 2],
+        "rate_code": [1, 2],
+        "total_fare": [100.0, 150.0],
+        "total_tips": [10.0, 15.0],
+        "total_amount": [110.0, 165.0],
+        "trip_count": [1, 2],
+    }
+)
+
 
 @patch("example.infrastructure.big_query.query")
 def test_get_trips(mock_big_query_query):
     any_start_date = date(2015, 1, 1)
     any_end_date = date(2015, 1, 2)
     any_payment_type = "Credit card"
 
-    trips = pd.DataFrame(
-        {
-            "day": ["2023-01-01"],
-            "payment_type": [1],
-            "rate_code": [1],
-            "total_fare": [100.0],
-            "total_tips": [10.0],
-            "total_amount": [110.0],
-            "trip_count": [1],
-        }
-    )
-    mock_big_query_query.return_value = trips
+    mock_big_query_query.return_value = _any_trips
 
     results = nyc_tlc_trips.get_trips((any_start_date, any_end_date), any_payment_type)
 
@@ -37,14 +38,21 @@ def test_get_trips(mock_big_query_query):
 
     expected_trips = pd.DataFrame(
         {
-            "day": ["2023-01-01"],
-            "payment_type": "Credit card",
-            "rate_code": "Standard rate",
-            "total_fare": [100.0],
-            "total_tips": [10.0],
-            "total_amount": [110.0],
-            "trip_count": [1],
+            "day": ["2023-01-01", "2023-01-02"],
+            "payment_type": ["Credit card", "Cash"],
+            "rate_code": ["Standard rate", "JFK"],
+            "total_fare": [100.0, 150.0],
+            "total_tips": [10.0, 15.0],
+            "total_amount": [110.0, 165.0],
+            "trip_count": [1, 2],
         }
     )
 
     pd.testing.assert_frame_equal(results, expected_trips)
+
+
+def test_export_trips():
+    csv = nyc_tlc_trips.export_trips(_any_trips).decode("utf-8")
+    lines = csv.split("\n")
+
+    assert lines[0] == """day,payment_type,rate_code,total_fare,total_tips,total_amount,trip_count"""