implement save function in argo

icepyx/quest/dataset_scripts/argo.py

@@ -1,3 +1,5 @@
+import os.path
+
 import numpy as np
 import pandas as pd
 import requests
@@ -404,7 +406,7 @@ def _parse_into_df(self, profile_data) -> pd.DataFrame:
 
         return profileDf
 
-    def download(self, params=None, presRange=None, keep_existing=True, savename='') -> pd.DataFrame:
+    def download(self, params=None, presRange=None, keep_existing=True) -> pd.DataFrame:
         """
         Downloads the requested data for a list of profile IDs (stored under .prof_ids) and returns it in a DataFrame.
 
@@ -487,6 +489,26 @@ def download(self, params=None, presRange=None, keep_existing=True, savename='')
 
         self.argodata.reset_index(inplace=True, drop=True)
 
-        if savename:
-            self.argodata.to_csv(savename + '_argo.csv')
         return self.argodata
+
+    def save(self, filepath):
+        """
+        Saves the argo dataframe to a csv at the specified location
+
+        Parameters
+        ----------
+        filepath: string containing complete filepath and name of file
+            extension will be removed and replaced with csv. Also appends
+            '_argo.csv' to filename
+            e.g. /path/to/file/my_data(.csv)
+        """
+
+
+        # create the directory if it doesn't exist
+        path, file = os.path.split(filepath)
+        if not os.path.exists(path):
+            os.mkdir(path)
+
+        # remove file extension
+        base, ext = os.path.splitext(filepath)
+        self.argodata.to_csv(base + '_argo.csv')

icepyx/quest/dataset_scripts/dataset.py

@@ -69,6 +69,12 @@ def download(self):
         """
         raise NotImplementedError
 
+    def save(self, filepath):
+        """
+        Save the downloaded data to a directory on your local machine.
+        """
+        raise NotImplementedError
+
     # ----------------------------------------------------------------------
     # Working with Data
 

icepyx/quest/quest.py

@@ -235,3 +235,13 @@ def download_all(self, path="", **kwargs):
             except:
                 dataset_name = type(v).__name__
                 print("Error downloading data from {0}".format(dataset_name))
+
+
+    def save_all(self, path):
+
+        for k, v in self.datasets.items():
+            if isinstance(v, Query):
+                print("ICESat-2 granules are saved during download")
+            else:
+                print("Saving " + k)
+                v.save(path)

icepyx/tests/test_quest_argo.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 import re
 
@@ -92,9 +94,9 @@ def test_presRange_setter_invalid_inputs(argo_quest_instance):
     reg_a = argo_quest_instance([-154, 30, -143, 37], ["2022-04-12", "2022-04-26"])
 
     exp = None
-    assert reg_a._presRange == exp
+    assert reg_a.presRange == exp
 
-    reg_a._presRange = (
+    reg_a.presRange = (
         "0.5, sam"  # it looks like the API will take a string with a space
     )
 
@@ -158,6 +160,16 @@ def test_download_parse_into_df(argo_quest_instance):
 # approach for additional testing of df functions: create json files with profiles and store them in test suite
 # then use those for the comparison (e.g. number of rows in df and json match)
 
+def test_save_df_to_csv(argo_quest_instance):
+    reg_a = argo_quest_instance([-154, 30, -143, 37], ["2022-04-12", "2022-04-13"])
+    reg_a.download()  # note: pressure is returned by default
+
+
+    path = os.getcwd() + "test_file"
+    reg_a.save(path)
+
+    assert os.path.exists(path + "_argo.csv")
+    os.remove(path + "_argo.csv")
 
 def test_merge_df(argo_quest_instance):
     reg_a = argo_quest_instance([-150, 30, -120, 60], ["2022-06-07", "2022-06-14"])
@@ -175,10 +187,11 @@ def test_merge_df(argo_quest_instance):
     assert "down_irradiance412_argoqc" in df.columns
 
 
+
+
 # ---------------------------------------------------
 # Test kwargs to replace params and presRange in search and download
 
-
 def test_replace_param_search(argo_quest_instance):
     reg_a = argo_quest_instance([-154, 30, -143, 37], ["2022-04-12", "2022-04-26"])