Update imf.py

bblocks/import_tools/imf.py

@@ -1,64 +1,34 @@
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from typing import Optional
 
 import pandas as pd
-import weo.dates
-from weo import all_releases, download, WEO
 
 from bblocks import config
 from bblocks.cleaning_tools.clean import clean_numeric_series, convert_to_datetime
 from bblocks.import_tools.common import ImportData
+from bblocks.import_tools.imf_weo import WEO
 from bblocks.logger import logger
 
 
-def _check_weo_parameters(
-    latest_y: int | None = None, latest_r: int | None = None
-) -> (int, int):
-    """Check parameters and return max values or provided input"""
-    if latest_y is None:
-        latest_y = max(*all_releases())[0]
+def _check_parameters(latest_y: int | None, latest_r: int | None) -> str | tuple:
+    if latest_y is None and latest_r is None:
+        release = "latest"
 
-    # if latest release isn't provided, take max value
-    if latest_r is None:
-        latest_r = max(*all_releases())[1]
+    else:
+        release = (latest_y, latest_r)
 
-    return latest_y, latest_r
+    return release
 
 
-def _update_weo(
-    latest_y: int = None,
-    latest_r: int = None,
-) -> None:
+def _update_weo(latest_y: int = None, latest_r: int = None) -> None:
     """Update _data from the World Economic Outlook, using WEO package"""
 
-    latest_y, latest_r = _check_weo_parameters(latest_y, latest_r)
+    release = _check_parameters(latest_y, latest_r)
 
     # Download the file from the IMF website and store in directory
-    download(
-        latest_y,
-        latest_r,
-        directory=config.BBPaths.raw_data,
-        filename=f"weo{latest_y}_{latest_r}.csv",
-    )
-
-    # Validate the file
-    if (
-        os.path.getsize(config.BBPaths.raw_data / f"weo{latest_y}_{latest_r}.csv")
-        < 1000
-    ):
-        print(
-            f"Downloading release {latest_r} of "
-            f"{latest_y} failed. Trying previous release"
-        )
-        os.remove(config.BBPaths.raw_data / f"weo{latest_y}_{latest_r}.csv")
-
-        try:
-            _update_weo(latest_y, latest_r - 1)
-        except weo.dates.DateError:
-            _update_weo(latest_y - 1, latest_r)
+    WEO(release).update_data()
 
 
 @dataclass
@@ -83,13 +53,11 @@ def __load_data(
             released value (1 or 2).
         """
 
-        latest_y, latest_r = _check_weo_parameters(latest_y, latest_r)
-
         names = {
             "ISO": "iso_code",
             "WEO Subject Code": "indicator",
             "Subject Descriptor": "indicator_name",
-            "Subject Notes": "indicator_description",
+            # "Country/Series-specific Notes": "indicator_description",
             "Units": "units",
             "Scale": "scale",
             "Estimates Start After": "estimates_start_after",
@@ -101,35 +69,17 @@ def __load_data(
             "Country/Series-specific Notes",
         ]
 
-        # If _data doesn't exist or update is required, update the _data
-        if not (config.BBPaths.raw_data / f"weo{latest_y}_{latest_r}.csv").exists():
-            _update_weo(latest_y, latest_r)
-
-        # Load the _data from disk. If it doesn't exist, try the previous one
-        try:
-            df = WEO(config.BBPaths.raw_data / f"weo{latest_y}_{latest_r}.csv").df
-            self.version = {"year": latest_y, "release": latest_r}
-        except FileNotFoundError:
-            try:
-                df = WEO(
-                    config.BBPaths.raw_data / f"weo{latest_y}_{latest_r - 1}.csv"
-                ).df
-                self.version = {"year": latest_y, "release": latest_r - 1}
-            except FileNotFoundError:
-                df = WEO(
-                    config.BBPaths.raw_data / f"weo{latest_y - 1}_{latest_r}.csv"
-                ).df
-                self.version = {"year": latest_y - 1, "release": latest_r}
+        release = _check_parameters(latest_y, latest_r)
+        df = WEO(version=release).load_data().get_old_format_data()
 
         # Load _data into _data object
         self._raw_data = (
-            df.drop(to_drop, axis=1)
+            df.drop(
+                columns=to_drop,
+            )
             .rename(columns=names)
             .melt(id_vars=names.values(), var_name="year", value_name="value")
-            .assign(
-                year=lambda d: convert_to_datetime(d.year),
-                value=lambda d: clean_numeric_series(d.value),
-            )
+            .assign(year=lambda d: convert_to_datetime(d.year))
             .dropna(subset=["value"])
             .reset_index(drop=True)
         )