Skip to content
Snippets Groups Projects
Commit bba861ec authored by Hugh Dickinson's avatar Hugh Dickinson
Browse files

Optional conversion to dataframe

parent 02e3722a
No related branches found
No related tags found
No related merge requests found
...@@ -34,7 +34,7 @@ class zooniverse: ...@@ -34,7 +34,7 @@ class zooniverse:
except PanoptesAPIException as e: except PanoptesAPIException as e:
return False 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...") print("Generating requested export...")
if wait: if wait:
print("\t\tWaiting for generation to complete...") print("\t\tWaiting for generation to complete...")
...@@ -44,6 +44,7 @@ class zooniverse: ...@@ -44,6 +44,7 @@ class zooniverse:
self._get_item_entry(item, "category"), generate=True, wait=wait self._get_item_entry(item, "category"), generate=True, wait=wait
) )
if response.ok: if response.ok:
if convertToFrame:
return ( return (
pd.read_csv( pd.read_csv(
io.BytesIO(response.content), io.BytesIO(response.content),
...@@ -54,10 +55,12 @@ class zooniverse: ...@@ -54,10 +55,12 @@ class zooniverse:
if wait if wait
else response else response
) )
else:
return response
else: else:
return None 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: if self.is_available(item) and not generate:
response = self._get_entity(item).get_export( response = self._get_entity(item).get_export(
self._get_item_entry(item, "category"), generate=False, wait=wait self._get_item_entry(item, "category"), generate=False, wait=wait
...@@ -78,16 +81,19 @@ class zooniverse: ...@@ -78,16 +81,19 @@ class zooniverse:
self._get_item_entry(item, "category"), generate=True, wait=wait self._get_item_entry(item, "category"), generate=True, wait=wait
) )
if response.ok: if response.ok:
if convertToFrame:
return ( return (
pd.read_csv( pd.read_csv(
io.BytesIO(response.content), io.BytesIO(response.content),
converters=zooniverse.category_converters[ converters=zooniverse.category_converters[
self._get_item_entry(item, "category") self._get_item_entry(item, "category")
], **read_csv_args ],
) )
if wait if wait
else response else response
) )
else:
return response
else: else:
return None return None
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment