Enhance type annotations for improved clarity and type safety in SGS and utils modules

Issue #26

Author: wilsonfreitas

bcb/sgs/__init__.py

@@ -1,10 +1,11 @@
 import json
 from io import StringIO
+from typing import Dict, Generator, List, Optional, Tuple, TypeAlias, Union
 
 import pandas as pd
 import requests
 
-from bcb.utils import Date
+from bcb.utils import Date, DateInput
 
 """
 Sistema Gerenciador de Séries Temporais (SGS)
@@ -17,7 +18,7 @@
 
 
 class SGSCode:
-    def __init__(self, code, name=None):
+    def __init__(self, code: Union[str, int], name: Optional[str] = None) -> None:
         if name is None:
             if isinstance(code, int) or isinstance(code, str):
                 self.name = str(code)
@@ -30,7 +31,16 @@ def __repr__(self):
         return f"{self.code} - {self.name}" if self.name else f"{self.code}"
 
 
-def _codes(codes):
+SGSCodeInput: TypeAlias = Union[
+    int,
+    str,
+    Tuple[str, Union[int, str]],
+    List[Union[int, str, Tuple[str, Union[int, str]]]],
+    Dict[str, Union[int, str]],
+]
+
+
+def _codes(codes: SGSCodeInput) -> Generator[SGSCode, None, None]:
     if isinstance(codes, int) or isinstance(codes, str):
         yield SGSCode(codes)
     elif isinstance(codes, tuple):
@@ -44,7 +54,7 @@ def _codes(codes):
             yield SGSCode(code, name)
 
 
-def _get_url_and_payload(code, start_date, end_date, last):
+def _get_url_and_payload(code: int, start_date: DateInput, end_date: DateInput, last: int) -> Dict[str, str]:
     payload = {"formato": "json"}
     if last == 0:
         if start_date is not None or end_date is not None:
@@ -58,7 +68,7 @@ def _get_url_and_payload(code, start_date, end_date, last):
     return {"payload": payload, "url": url}
 
 
-def _format_df(df, code, freq):
+def _format_df(df: pd.DataFrame, code: SGSCode, freq: str) -> pd.DataFrame:
     cns = {"data": "Date", "valor": code.name, "datafim": "enddate"}
     df = df.rename(columns=cns)
     if "Date" in df:
@@ -71,7 +81,14 @@ def _format_df(df, code, freq):
     return df
 
 
-def get(codes, start=None, end=None, last=0, multi=True, freq=None):
+def get(
+    codes: SGSCodeInput,
+    start: Optional[DateInput] = None,
+    end: Optional[DateInput] = None,
+    last: int = 0,
+    multi: bool = True,
+    freq: Optional[str] = None,
+) -> Union[pd.DataFrame, List[pd.DataFrame]]:
     """
     Retorna um DataFrame pandas com séries temporais obtidas do SGS.
 
@@ -130,7 +147,7 @@ def get(codes, start=None, end=None, last=0, multi=True, freq=None):
             return dfs
 
 
-def get_json(code: int, start=None, end=None, last: int = 0) -> str:
+def get_json(code: int, start: Optional[DateInput] = None, end: Optional[DateInput] = None, last: int = 0) -> str:
     """
     Retorna um JSON com séries temporais obtidas do SGS.
 

bcb/sgs/regional_economy.py

@@ -1,4 +1,9 @@
+from typing import Dict, List, Optional, Union
+
+import pandas as pd
+
 from bcb.sgs import get
+from bcb.utils import DateInput
 
 """
 Dados da Economia Regional
@@ -119,7 +124,7 @@
 }
 
 
-def get_non_performing_loans_codes(states_or_region, mode="total"):
+def get_non_performing_loans_codes(states_or_region: Union[str, List[str]], mode: str = "total") -> Dict[str, str]:
     is_state = False
     is_region = False
     states_or_region = [states_or_region] if isinstance(states_or_region, str) else states_or_region
@@ -151,7 +156,14 @@ def get_non_performing_loans_codes(states_or_region, mode="total"):
     return codes
 
 
-def get_non_performing_loans(states_or_region, mode="total", start=None, end=None, last=0, freq=None):
+def get_non_performing_loans(
+    states_or_region: Union[str, List[str]],
+    mode: str = "total",
+    start: Optional[DateInput] = None,
+    end: Optional[DateInput] = None,
+    last: int = 0,
+    freq: Optional[str] = None,
+) -> pd.DataFrame:
     """Dados de inadimplência das operações de crédito.
 
     Esta função é um *wrapper* para o método para a função ``get`` do módulo ``sgs``,

bcb/utils.py

@@ -1,5 +1,5 @@
 from datetime import date, datetime
-from typing import Union
+from typing import TypeAlias, Union
 
 BRAZILIAN_REGIONS = {
     "N": ["AC", "AP", "AM", "PA", "RO", "RR", "TO"],
@@ -12,9 +12,11 @@
 for state in BRAZILIAN_REGIONS.values():
     BRAZILIAN_STATES.extend(state)
 
+DateInput: TypeAlias = Union[str, datetime, "Date", date]
+
 
 class Date:
-    def __init__(self, d: Union[str, datetime, "Date", date], format: str = "%Y-%m-%d") -> None:
+    def __init__(self, d: DateInput, format: str = "%Y-%m-%d") -> None:
         if isinstance(d, str):
             if d == "now" or d == "today":
                 d = date.today()