Defined config dataclass structure (#142)

* Defined config dataclass structure.

* fixed flake8 errors

* made changes in the docstring and structure of config dataclass

* implemented config-dataclass in weather-dl

* updated test cases after implementation of config-dataclass

* moved creation of config-dataclass object in process_config method and few small changes

weather_dl/download_pipeline/clients.py

@@ -26,6 +26,7 @@
 
 import cdsapi
 from ecmwfapi import ECMWFService
+from .config import Config
 
 warnings.simplefilter(
     "ignore", category=urllib3.connectionpool.InsecureRequestWarning)
@@ -42,7 +43,7 @@ class Client(abc.ABC):
         level: Default log level for the client.
     """
 
-    def __init__(self, config: t.Dict, level: int = logging.INFO) -> None:
+    def __init__(self, config: Config, level: int = logging.INFO) -> None:
         """Clients are initialized with the general CLI configuration."""
         self.config = config
         self.logger = logging.getLogger(f'{__name__}.{type(self).__name__}')
@@ -87,11 +88,11 @@ class CdsClient(Client):
     """Name patterns of datasets that are hosted internally on CDS servers."""
     cds_hosted_datasets = {'reanalysis-era'}
 
-    def __init__(self, config: t.Dict, level: int = logging.INFO) -> None:
+    def __init__(self, config: Config, level: int = logging.INFO) -> None:
         super().__init__(config, level)
         self.c = cdsapi.Client(
-            url=config['parameters'].get('api_url', os.environ.get('CDSAPI_URL')),
-            key=config['parameters'].get('api_key', os.environ.get('CDSAPI_KEY')),
+            url=config.kwargs.get('api_url', os.environ.get('CDSAPI_URL')),
+            key=config.kwargs.get('api_key', os.environ.get('CDSAPI_KEY')),
             debug_callback=self.logger.debug,
             info_callback=self.logger.info,
             warning_callback=self.logger.warning,
@@ -171,13 +172,13 @@ class MarsClient(Client):
         level: Default log level for the client.
     """
 
-    def __init__(self, config: t.Dict, level: int = logging.INFO) -> None:
+    def __init__(self, config: Config, level: int = logging.INFO) -> None:
         super().__init__(config, level)
         self.c = ECMWFService(
             "mars",
-            key=config['parameters'].get('api_key', os.environ.get("MARSAPI_KEY")),
-            url=config['parameters'].get('api_url', os.environ.get("MARSAPI_URL")),
-            email=config['parameters'].get('api_email', os.environ.get("MARSAPI_EMAIL")),
+            key=config.kwargs.get('api_key', os.environ.get("MARSAPI_KEY")),
+            url=config.kwargs.get('api_url', os.environ.get("MARSAPI_URL")),
+            email=config.kwargs.get('api_email', os.environ.get("MARSAPI_EMAIL")),
             log=self.logger.debug,
             verbose=True
         )

weather_dl/download_pipeline/clients_test.py

@@ -15,26 +15,27 @@
 import unittest
 
 from .clients import FakeClient, CdsClient, MarsClient
+from .config import Config
 
 
 class MaxWorkersTest(unittest.TestCase):
     def test_cdsclient_internal(self):
-        client = CdsClient({'parameters': {'api_url': 'url', 'api_key': 'key'}})
+        client = CdsClient(Config.from_dict({'parameters': {'api_url': 'url', 'api_key': 'key'}}))
         self.assertEqual(
             client.num_requests_per_key("reanalysis-era5-some-data"), 5)
 
     def test_cdsclient_mars_hosted(self):
-        client = CdsClient({'parameters': {'api_url': 'url', 'api_key': 'key'}})
+        client = CdsClient(Config.from_dict({'parameters': {'api_url': 'url', 'api_key': 'key'}}))
         self.assertEqual(
             client.num_requests_per_key("reanalysis-carra-height-levels"), 2)
 
     def test_marsclient(self):
-        client = MarsClient({'parameters': {}})
+        client = MarsClient(Config.from_dict({'parameters': {}}))
         self.assertEqual(
             client.num_requests_per_key("reanalysis-era5-some-data"), 2)
 
     def test_fakeclient(self):
-        client = FakeClient({'parameters': {}})
+        client = FakeClient(Config.from_dict({'parameters': {}}))
         self.assertEqual(
             client.num_requests_per_key("reanalysis-era5-some-data"), 1)
 

weather_dl/download_pipeline/config.py

@@ -0,0 +1,71 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import typing as t
+import dataclasses
+
+Values = t.Union[t.List['Values'], t.Dict[str, 'Values'], bool, int, float, str]  # pytype: disable=not-supported-yet
+
+
+@dataclasses.dataclass
+class Config:
+    """Contains pipeline parameters.
+
+    Attributes:
+        client:
+            Name of the Weather-API-client. Supported clients are mentioned in the 'CLIENTS' variable.
+        dataset (optional):
+            Name of the target dataset. Allowed options are dictated by the client.
+        partition_keys (optional):
+            Choose the keys from the selection section to partition the data request.
+            This will compute a cartesian cross product of the selected keys
+            and assign each as their own download.
+        target_path:
+            Download artifact filename template. Can make use of Python's standard string formatting.
+            It can contain format symbols to be replaced by partition keys;
+            if this is used, the total number of format symbols must match the number of partition keys.
+        subsection_name:
+            Name of the particular subsection. 'default' if there is no subsection.
+        force_download:
+            Force redownload of partitions that were previously downloaded.
+        user_id:
+            Username from the environment variables.
+        kwargs (optional):
+            For representing subsections or any other parameters.
+        selection:
+            Contains parameters used to select desired data.
+    """
+
+    client: str = ""
+    dataset: t.Optional[str] = ""
+    target_path: str = ""
+    partition_keys: t.Optional[t.List[str]] = dataclasses.field(default_factory=list)
+    subsection_name: str = "default"
+    force_download: bool = False
+    user_id: str = "unknown"
+    kwargs: t.Optional[t.Dict[str, Values]] = dataclasses.field(default_factory=dict)
+    selection: t.Dict[str, Values] = dataclasses.field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, config: t.Dict) -> 'Config':
+        config_instance = cls()
+        for section_key, section_value in config.items():
+            if section_key == "parameters":
+                for key, value in section_value.items():
+                    if hasattr(config_instance, key):
+                        setattr(config_instance, key, value)
+                    else:
+                        config_instance.kwargs[key] = value
+            if section_key == "selection":
+                config_instance.selection = section_value
+        return config_instance

weather_dl/download_pipeline/fetcher.py

@@ -23,6 +23,7 @@
 from .clients import CLIENTS, Client
 from .manifest import Manifest, NoOpManifest, Location
 from .parsers import prepare_target_name
+from .config import Config
 from .partition import skip_partition
 from .stores import Store, FSStore
 from .util import retry_with_exponential_backoff
@@ -64,7 +65,7 @@ def retrieve(self, client: Client, dataset: str, selection: t.Dict, dest: str) -
         """Retrieve from download client, with retries."""
         client.retrieve(dataset, selection, dest)
 
-    def fetch_data(self, config: t.Dict, *, worker_name: str = 'default') -> None:
+    def fetch_data(self, config: Config, *, worker_name: str = 'default') -> None:
         """Download data from a client to a temp file, then upload to Cloud Storage."""
         if not config:
             return
@@ -74,14 +75,11 @@ def fetch_data(self, config: t.Dict, *, worker_name: str = 'default') -> None:
 
         client = CLIENTS[self.client_name](config)
         target = prepare_target_name(config)
-        dataset = config['parameters'].get('dataset', '')
-        selection = config['selection']
-        user = config['parameters'].get('user_id', 'unknown')
 
-        with self.manifest.transact(selection, target, user):
+        with self.manifest.transact(config.selection, target, config.user_id):
             with tempfile.NamedTemporaryFile() as temp:
                 logger.info(f'[{worker_name}] Fetching data for {target!r}.')
-                self.retrieve(client, dataset, selection, temp.name)
+                self.retrieve(client, config.dataset, config.selection, temp.name)
 
                 logger.info(f'[{worker_name}] Uploading to store for {target!r}.')
                 self.upload(temp, target)

weather_dl/download_pipeline/fetcher_test.py

@@ -22,6 +22,7 @@
 from .fetcher import Fetcher
 from .manifest import MockManifest, Location
 from .stores import InMemoryStore, FSStore
+from .config import Config
 
 
 class UploadTest(unittest.TestCase):
@@ -69,7 +70,7 @@ def setUp(self) -> None:
     @patch('weather_dl.download_pipeline.stores.InMemoryStore.open', return_value=io.StringIO())
     @patch('cdsapi.Client.retrieve')
     def test_fetch_data(self, mock_retrieve, mock_gcs_file):
-        config = {
+        config = Config.from_dict({
             'parameters': {
                 'dataset': 'reanalysis-era5-pressure-levels',
                 'partition_keys': ['year', 'month'],
@@ -82,7 +83,7 @@ def test_fetch_data(self, mock_retrieve, mock_gcs_file):
                 'month': ['12'],
                 'year': ['01']
             }
-        }
+        })
 
         fetcher = Fetcher('cds', self.dummy_manifest, InMemoryStore())
         fetcher.fetch_data(config)
@@ -94,13 +95,13 @@ def test_fetch_data(self, mock_retrieve, mock_gcs_file):
 
         mock_retrieve.assert_called_with(
             'reanalysis-era5-pressure-levels',
-            config['selection'],
+            config.selection,
             ANY)
 
     @patch('weather_dl.download_pipeline.stores.InMemoryStore.open', return_value=io.StringIO())
     @patch('cdsapi.Client.retrieve')
     def test_fetch_data__manifest__returns_success(self, mock_retrieve, mock_gcs_file):
-        config = {
+        config = Config.from_dict({
             'parameters': {
                 'dataset': 'reanalysis-era5-pressure-levels',
                 'partition_keys': ['year', 'month'],
@@ -113,13 +114,13 @@ def test_fetch_data__manifest__returns_success(self, mock_retrieve, mock_gcs_fil
                 'month': ['12'],
                 'year': ['01']
             }
-        }
+        })
 
         fetcher = Fetcher('cds', self.dummy_manifest, InMemoryStore())
         fetcher.fetch_data(config)
 
         self.assertDictContainsSubset(dict(
-            selection=config['selection'],
+            selection=config.selection,
             location='gs://weather-dl-unittest/download-01-12.nc',
             status='success',
             error=None,
@@ -128,7 +129,7 @@ def test_fetch_data__manifest__returns_success(self, mock_retrieve, mock_gcs_fil
 
     @patch('cdsapi.Client.retrieve')
     def test_fetch_data__manifest__records_retrieve_failure(self, mock_retrieve):
-        config = {
+        config = Config.from_dict({
             'parameters': {
                 'dataset': 'reanalysis-era5-pressure-levels',
                 'partition_keys': ['year', 'month'],
@@ -141,7 +142,7 @@ def test_fetch_data__manifest__records_retrieve_failure(self, mock_retrieve):
                 'month': ['12'],
                 'year': ['01']
             }
-        }
+        })
 
         error = IOError("We don't have enough permissions to download this.")
         mock_retrieve.side_effect = error
@@ -153,7 +154,7 @@ def test_fetch_data__manifest__records_retrieve_failure(self, mock_retrieve):
         actual = list(self.dummy_manifest.records.values())[0]._asdict()
 
         self.assertDictContainsSubset(dict(
-            selection=config['selection'],
+            selection=config.selection,
             location='gs://weather-dl-unittest/download-01-12.nc',
             status='failure',
             user='unknown',
@@ -165,7 +166,7 @@ def test_fetch_data__manifest__records_retrieve_failure(self, mock_retrieve):
     @patch('weather_dl.download_pipeline.stores.InMemoryStore.open', return_value=io.StringIO())
     @patch('cdsapi.Client.retrieve')
     def test_fetch_data__manifest__records_gcs_failure(self, mock_retrieve, mock_gcs_file):
-        config = {
+        config = Config.from_dict({
             'parameters': {
                 'dataset': 'reanalysis-era5-pressure-levels',
                 'partition_keys': ['year', 'month'],
@@ -178,7 +179,7 @@ def test_fetch_data__manifest__records_gcs_failure(self, mock_retrieve, mock_gcs
                 'month': ['12'],
                 'year': ['01']
             }
-        }
+        })
 
         error = IOError("Can't open gcs file.")
         mock_gcs_file.side_effect = error
@@ -189,7 +190,7 @@ def test_fetch_data__manifest__records_gcs_failure(self, mock_retrieve, mock_gcs
 
         actual = list(self.dummy_manifest.records.values())[0]._asdict()
         self.assertDictContainsSubset(dict(
-            selection=config['selection'],
+            selection=config.selection,
             location='gs://weather-dl-unittest/download-01-12.nc',
             status='failure',
             user='unknown',
@@ -201,7 +202,7 @@ def test_fetch_data__manifest__records_gcs_failure(self, mock_retrieve, mock_gcs
     @patch('weather_dl.download_pipeline.stores.InMemoryStore.open', return_value=io.StringIO())
     @patch('cdsapi.Client.retrieve')
     def test_fetch_data__skips_existing_download(self, mock_retrieve, mock_gcs_file):
-        config = {
+        config = Config.from_dict({
             'parameters': {
                 'dataset': 'reanalysis-era5-pressure-levels',
                 'partition_keys': ['year', 'month'],
@@ -214,7 +215,7 @@ def test_fetch_data__skips_existing_download(self, mock_retrieve, mock_gcs_file)
                 'month': ['12'],
                 'year': ['01']
             }
-        }
+        })
 
         # target file already exists in store...
         store = InMemoryStore()