Skip to content

Commit

Permalink
Feature/source ads (#142)
Browse files Browse the repository at this point in the history
* add preliminary ADS source
* add simple tests
* refactor to more closely follow climetlab version of ads.py
  • Loading branch information
TomHodson authored Aug 7, 2023
1 parent 030b9eb commit fdae1af
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 2 deletions.
75 changes: 75 additions & 0 deletions earthkit/data/sources/ads.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
# (C) Copyright 2020 ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.
#

import os

import cdsapi
import yaml

from .cds import CdsRetriever
from .prompt import APIKeyPrompt


class ADSAPIKeyPrompt(APIKeyPrompt):
register_or_sign_in_url = "https://ads.atmosphere.copernicus.eu/"
retrieve_api_key_url = "https://ads.atmosphere.copernicus.eu/api-how-to"

prompts = [
dict(
name="url",
default="https://ads.atmosphere.copernicus.eu/api/v2",
title="API url",
validate=r"http.?://.*",
),
dict(
name="key",
example="123:abcdef01-0000-1111-2222-0123456789ab",
title="API key",
hidden=True,
validate=r"\d+:[\-0-9a-f]+",
),
]

rcfile = "~/.adsapirc"

def save(self, input, file):
yaml.dump(input, file, default_flow_style=False)


def client():
prompt = ADSAPIKeyPrompt()
prompt.check()

path = os.path.expanduser(prompt.rcfile)

if not os.path.exists(path):
prompt.ask_user_and_save()

with open(path) as f:
rc = yaml.safe_load(f.read())

return cdsapi.Client(**rc)


EXTENSIONS = {
"grib": ".grib",
"netcdf": ".nc",
}


class AdsRetriever(CdsRetriever):
sphinxdoc = """
AdsRetriever
"""

def client(self):
return client()


source = AdsRetriever
7 changes: 5 additions & 2 deletions earthkit/data/sources/cds.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,9 @@ class CdsRetriever(FileSource):
CdsRetriever
"""

def client(self):
return client()

def __init__(self, dataset, *args, **kwargs):
super().__init__()

Expand All @@ -81,7 +84,7 @@ def __init__(self, dataset, *args, **kwargs):

requests = self.requests(**kwargs)

client() # Trigger password prompt before thraeding
self.client() # Trigger password prompt before thraeding

nthreads = min(self.settings("number-of-download-threads"), len(requests))

Expand All @@ -96,7 +99,7 @@ def __init__(self, dataset, *args, **kwargs):

def _retrieve(self, dataset, request):
def retrieve(target, args):
client().retrieve(args[0], args[1], target)
self.client().retrieve(args[0], args[1], target)

return self.cache_file(
retrieve,
Expand Down
86 changes: 86 additions & 0 deletions tests/sources/test_ads.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
#!/usr/bin/env python3

# (C) Copyright 2020 ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.
#

import os

import pytest

from earthkit.data import from_source

NO_ADS = not os.path.exists(os.path.expanduser("~/.adsapirc"))


@pytest.mark.long_test
@pytest.mark.download
@pytest.mark.skipif(NO_ADS, reason="No access to ADS")
def test_ads_grib_1():
s = from_source(
"ads",
"cams-global-reanalysis-eac4",
variable=["particulate_matter_10um", "particulate_matter_1um"],
area=[50, -50, 20, 50],
date="2012-12-12",
time="12:00",
)
assert len(s) == 2


@pytest.mark.long_test
@pytest.mark.download
@pytest.mark.skipif(NO_ADS, reason="No access to ADS")
def test_ads_grib_2():
s = from_source(
"ads",
"cams-global-reanalysis-eac4",
variable=["particulate_matter_10um", "particulate_matter_1um"],
area=[50, -50, 20, 50],
date="2012-12-12",
time="12:00",
split_on="variable",
)
assert len(s) == 2


@pytest.mark.long_test
@pytest.mark.download
@pytest.mark.skipif(NO_ADS, reason="No access to ADS")
def test_ads_grib_3():
s = from_source(
"ads",
"cams-global-reanalysis-eac4",
variable=["particulate_matter_10um", "particulate_matter_1um"],
area=[50, -50, 20, 50],
date="2012-12-12/to/2012-12-15",
time="12:00",
)
assert len(s) == 8


@pytest.mark.long_test
@pytest.mark.download
@pytest.mark.skipif(NO_ADS, reason="No access to ADS")
def test_ads_netcdf():
s = from_source(
"ads",
"cams-global-reanalysis-eac4",
variable=["particulate_matter_10um", "particulate_matter_1um"],
area=[50, -50, 20, 50],
date="2012-12-12",
time="12:00",
format="netcdf",
)
assert len(s) == 2


if __name__ == "__main__":
from earthkit.data.testing import main

main(__file__)

0 comments on commit fdae1af

Please sign in to comment.