From bba861ec17dba797d04ac3c9feaad0b6228fd9bb Mon Sep 17 00:00:00 2001
From: Hugh Dickinson <hugh.dickinson@open.ac.uk>
Date: Thu, 24 Jun 2021 15:05:13 +0100
Subject: [PATCH] Optional conversion to dataframe

---
 zooniverse/zooniverse.py | 46 +++++++++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 20 deletions(-)

diff --git a/zooniverse/zooniverse.py b/zooniverse/zooniverse.py
index fb31142..0e8fcca 100644
--- a/zooniverse/zooniverse.py
+++ b/zooniverse/zooniverse.py
@@ -34,7 +34,7 @@ class zooniverse:
         except PanoptesAPIException as e:
             return False
 
-    def generate(self, item, wait=False, **read_csv_args):
+    def generate(self, item, wait=False, convertToFrame=True, **read_csv_args):
         print("Generating requested export...")
         if wait:
             print("\t\tWaiting for generation to complete...")
@@ -44,20 +44,23 @@ class zooniverse:
             self._get_item_entry(item, "category"), generate=True, wait=wait
         )
         if response.ok:
-            return (
-                pd.read_csv(
-                    io.BytesIO(response.content),
-                    converters=zooniverse.category_converters[
-                        self._get_item_entry(item, "category")
-                    ],
+            if convertToFrame:
+                return (
+                    pd.read_csv(
+                        io.BytesIO(response.content),
+                        converters=zooniverse.category_converters[
+                            self._get_item_entry(item, "category")
+                        ],
+                    )
+                    if wait
+                    else response
                 )
-                if wait
-                else response
-            )
+            else:
+                return response
         else:
             return None
 
-    def retrieve(self, item, generate=False, wait=False, **read_csv_args):
+    def retrieve(self, item, generate=False, wait=False, convertToFrame=True, **read_csv_args):
         if self.is_available(item) and not generate:
             response = self._get_entity(item).get_export(
                 self._get_item_entry(item, "category"), generate=False, wait=wait
@@ -78,16 +81,19 @@ class zooniverse:
                     self._get_item_entry(item, "category"), generate=True, wait=wait
                 )
         if response.ok:
-            return (
-                pd.read_csv(
-                    io.BytesIO(response.content),
-                    converters=zooniverse.category_converters[
-                        self._get_item_entry(item, "category")
-                    ], **read_csv_args
+            if convertToFrame:
+                return (
+                    pd.read_csv(
+                        io.BytesIO(response.content),
+                        converters=zooniverse.category_converters[
+                            self._get_item_entry(item, "category")
+                        ],
+                    )
+                    if wait
+                    else response
                 )
-                if wait
-                else response
-            )
+            else:
+                return response
         else:
             return None
 
-- 
GitLab