Merge pull request #141 from israelmcmc/wasabi_download

Add fetch_wasabi_file() utility function.

cosipy/util/__init__.py

@@ -0,0 +1 @@
+from .data_fetching import fetch_wasabi_file

cosipy/util/data_fetching.py

@@ -0,0 +1,46 @@
+import os
+from awscli.clidriver import create_clidriver
+
+def fetch_wasabi_file(file,
+                      output = None,
+                      override = False,
+                      bucket = 'cosi-pipeline-public',
+                      endpoint = 'https://s3.us-west-1.wasabisys.com',
+                      access_key = 'GBAL6XATQZNRV3GFH9Y4',
+                      secret_key = 'GToOczY5hGX3sketNO2fUwiq4DJoewzIgvTCHoOv'):
+    """
+    Download a file from COSI's Wasabi acccount.
+
+    Parameters
+    ----------
+    file : str
+        Full path to file in Wasabi
+    output : str,  optional
+        Full path to the downloaded file in the local system. By default it will use 
+        the current durectory and the same file name as the input file.
+    bucket : str, optional
+        Passed to aws --bucket option
+    endpoint : str, optional
+        Passed to aws --endpoint-url option
+    access_key : str, optional
+        AWS_ACCESS_KEY_ID
+    secret_key : str, optional
+        AWS_SECRET_ACCESS_KEY
+    """
+
+    if output is None:
+        output = file.split('/')[-1]
+
+    if os.path.exists(output) and not override:
+        raise RuntimeError(f"File {output} already exists.")
+
+    cli = create_clidriver()
+
+    cli.session.set_credentials(access_key, secret_key)
+
+    cli.main(['s3api', 'get-object',
+              '--bucket', bucket,
+              '--key', file,
+              '--endpoint-url', endpoint,
+              output])
+

docs/api/index.rst

@@ -18,6 +18,6 @@ If you are instead interested in an overview on how to use cosipy, see out `tuto
    threeml
    ts_map
    image_deconvolution
+   util          
 
 
-

docs/api/util.rst

@@ -0,0 +1,4 @@
+Utilities
+=========
+
+.. autofunction:: cosipy.util.fetch_wasabi_file

docs/conf.py

@@ -65,7 +65,8 @@
                         'scoords',
                         'pandas',
                         'tqdm',
-                        'scipy']
+                        'scipy',
+                        'awscli']
 
 # There seems to be a conflict between unittest.mock (used by sphinx) and metaclasses
 # The cosipy.threeml.custom_functions.Band_Eflux includes a metaclass from

tests/util/test_data_fetching.py

@@ -0,0 +1,8 @@
+from cosipy.util import fetch_wasabi_file
+
+fetch_wasabi_file('test_file.txt', override = True)
+
+f = open('test_file.txt')
+
+assert f.read() == 'Small file used for testing purposes.\n'
+