From 11de34cf413e3e9761eb6d276a88ed1149ea6519 Mon Sep 17 00:00:00 2001
From: chasem <chasem@wustl.edu>
Date: Sun, 7 Jul 2024 02:48:20 +0000
Subject: [PATCH] adjusting exportTableAsGzipFileMixin so that it actually
 exports as a gzip; fixes #41

---
 .../mixins/ExportTableAsGzipFileMixin.py      | 27 ++++++++++++-------
 .../regulatory_data/tests/test_filters.py     |  9 ++++---
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/yeastregulatorydb/regulatory_data/api/views/mixins/ExportTableAsGzipFileMixin.py b/yeastregulatorydb/regulatory_data/api/views/mixins/ExportTableAsGzipFileMixin.py
index f8f5805..d01b0fc 100644
--- a/yeastregulatorydb/regulatory_data/api/views/mixins/ExportTableAsGzipFileMixin.py
+++ b/yeastregulatorydb/regulatory_data/api/views/mixins/ExportTableAsGzipFileMixin.py
@@ -1,3 +1,6 @@
+import gzip
+from io import BytesIO
+
 import pandas as pd
 from django.http import HttpResponse
 from rest_framework.decorators import action
@@ -13,17 +16,23 @@ def export(self, request):
         # Get the queryset and apply any filters
         queryset = self.filter_queryset(self.get_queryset())
 
-        # Convert the filtered queryset to a DataFrame
-        df = pd.DataFrame.from_records(queryset.values())
+        # Serialize the queryset
+        serializer = self.get_serializer(queryset, many=True)
+        serialized_data = serializer.data
 
-        # Create a HttpResponse object with the appropriate CSV header.
-        response = HttpResponse(content_type="text/csv")
-        response["Content-Disposition"] = f'attachment; filename="{self.queryset.model.__name__}.csv.gz"'
+        # Convert the serialized data to a DataFrame
+        df = pd.DataFrame.from_records(serialized_data)
 
-        # Write the DataFrame to a CSV string and compress it
-        csv_data = df.to_csv(index=False, compression="gzip")
+        # Create a CSV string
+        csv_data = df.to_csv(index=False)
 
-        # Set the compressed CSV data as the content of the response
-        response.content = csv_data
+        # Compress the CSV data using gzip
+        gzip_buffer = BytesIO()
+        with gzip.GzipFile(fileobj=gzip_buffer, mode="w") as f:
+            f.write(csv_data.encode("utf-8"))
+
+        # Create a HttpResponse object with the appropriate CSV header.
+        response = HttpResponse(gzip_buffer.getvalue(), content_type="application/gzip")
+        response["Content-Disposition"] = f'attachment; filename="{self.queryset.model.__name__}.csv.gz"'
 
         return response
diff --git a/yeastregulatorydb/regulatory_data/tests/test_filters.py b/yeastregulatorydb/regulatory_data/tests/test_filters.py
index 4bcedb6..2849908 100644
--- a/yeastregulatorydb/regulatory_data/tests/test_filters.py
+++ b/yeastregulatorydb/regulatory_data/tests/test_filters.py
@@ -433,7 +433,7 @@ def test_genomic_feature_filter():
     filter_params = [
         {"chr": "chr1"},
         {"start_min": 1, "start_max": 50},  # Test range filter for start
-        {"end_min": 50, "end_max": 150},    # Test range filter for end
+        {"end_min": 50, "end_max": 150},  # Test range filter for end
         {"strand": "+"},
         {"type": "type1"},
         {"locus_tag": "tag1"},
@@ -452,7 +452,7 @@ def test_genomic_feature_filter():
     # Additional test cases for range filters
     range_filter_params = [
         {"start_min": 0, "start_max": 200},  # Both genomic_feature1 and genomic_feature2 should be included
-        {"end_min": 1, "end_max": 150},      # Only genomic_feature1 should be included
+        {"end_min": 1, "end_max": 150},  # Only genomic_feature1 should be included
     ]
 
     # Apply each range filter and check the expected results
@@ -460,7 +460,10 @@ def test_genomic_feature_filter():
         f = GenomicFeatureFilter(params, queryset=GenomicFeature.objects.all())
         if "start_min" in params or "start_max" in params:
             assert genomic_feature1 in f.qs, f"Failed for range filter params: {params}"
-            assert genomic_feature2 in f.qs, f"Failed for range filter params: {params}" if params["start_max"] == 200 else assert genomic_feature2 not in f.qs, f"Failed for range filter params: {params}"
+        if params.get("start_max", None) == 200:
+            assert genomic_feature2 in f.qs, f"Failed for range filter params: {params}"
+        else:
+            assert genomic_feature2 not in f.qs, f"Failed for range filter params: {params}"
         if "end_min" in params or "end_max" in params:
             assert genomic_feature1 in f.qs, f"Failed for range filter params: {params}"
             assert genomic_feature2 not in f.qs, f"Failed for range filter params: {params}"