fix: Resolve mypy errors

experimental/download_databases.py

@@ -2,13 +2,14 @@
 # recommend trying this.
 
 import imfp
+from pandas import DataFrame
 
 # Set a custom wait time
 _imf_wait_time = 10
 
 # Attempt to download databases sequentially
-databases = imfp.imf_databases()
-datasets = {"database_names": [], "dataframes": []}
+databases: DataFrame = imfp.imf_databases()
+datasets: dict[str, list[DataFrame | None]] = {"database_names": [], "dataframes": []}
 for database_id in databases["database_id"]:
     datasets["database_names"].append(database_id)
     try:

experimental/download_indicators.py

@@ -1,6 +1,7 @@
 # This script tries to download whole databases indicator by indicator
 
 import imfp
+from pandas import DataFrame
 
 # Examine database list
 databases = imfp.imf_databases()
@@ -11,8 +12,8 @@
 
 # Try to download the database indicator by indicator (Note that some databases don't
 # use the 'indicator' parameter, so this won't work with every database)
-indicators = imfp.imf_parameters("IFS")["indicator"]
-datasets = {"indicator_names": [], "dataframes": []}
+indicators: DataFrame = imfp.imf_parameters("IFS")["indicator"]
+datasets: dict[str, list[DataFrame | None]] = {"indicator_names": [], "dataframes": []}
 for indicator in indicators["input_code"]:
     datasets["indicator_names"].append(indicator)
     try:

imfp/data.py

@@ -1,8 +1,10 @@
 import logging
-from pandas import DataFrame, Series, concat
+from typing import overload, Literal
 from warnings import warn
-from .utils import _download_parse, _imf_dimensions, _imf_metadata
 from urllib.parse import urlencode
+from pandas import DataFrame, Series, concat
+
+from .utils import _download_parse, _imf_dimensions, _imf_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -175,11 +177,38 @@ def imf_parameter_defs(
     return parameterlist
 
 
+@overload
 def imf_dataset(
     database_id: str,
-    parameters: dict = None,
-    start_year: int = None,
-    end_year: int = None,
+    parameters: dict | None = None,
+    start_year: int | str | None = None,
+    end_year: int | str | None = None,
+    return_raw: bool = False,
+    print_url: bool = False,
+    times: int = 3,
+    include_metadata: Literal[False] = False,
+    **kwargs,
+) -> DataFrame: ...
+
+@overload
+def imf_dataset(
+    database_id: str,
+    parameters: dict | None = None,
+    start_year: int | str | None = None,
+    end_year: int | str | None = None,
+    return_raw: bool = False,
+    print_url: bool = False,
+    times: int = 3,
+    include_metadata: Literal[True] = True,
+    **kwargs,
+) -> tuple[DataFrame, DataFrame]: ...
+
+
+def imf_dataset(
+    database_id: str,
+    parameters: dict | None = None,
+    start_year: int | str | None = None,
+    end_year: int | str | None = None,
     return_raw: bool = False,
     print_url: bool = False,
     times: int = 3,
@@ -224,10 +253,7 @@ def imf_dataset(
         return_raw == True, returns the raw JSON fetched from the API endpoint.
     """
 
-    if database_id is None:
-        raise ValueError("Missing required database_id argument.")
-
-    if not isinstance(database_id, str):
+    if not database_id or not isinstance(database_id, str):
         raise ValueError("database_id must be a string.")
 
     years = {}