Add location as name

Author: anapaulagomes

bcb/sgs/__init__.py

@@ -25,6 +25,9 @@ def __init__(self, code, name=None):
             self.name = str(name)
             self.value = int(code)
 
+    def __repr__(self):
+        return f"{self.code} - {self.name}" if self.name else f"{self.code}"
+
 
 def _codes(codes):
     if isinstance(codes, int) or isinstance(codes, str):
@@ -36,8 +39,8 @@ def _codes(codes):
             _ist = isinstance(cd, tuple)
             yield SGSCode(cd[1], cd[0]) if _ist else SGSCode(cd)
     elif isinstance(codes, dict):
-        for cd in codes:
-            yield SGSCode(codes[cd], cd)
+        for name, code in codes.items():
+            yield SGSCode(code, name)
 
 
 def _get_url_and_payload(code, start_date, end_date, last):

bcb/sgs/regional_economy.py

@@ -127,7 +127,7 @@ def get_non_performing_loans_codes(states_or_region, mode="total"):
     if not is_state and not is_region:
         raise Exception(f"Not a valid state or region: {states_or_region}")
 
-    codes = []
+    codes = {}
     non_performing_loans_by_location = NON_PERFORMING_LOANS_BY_STATE_TOTAL
     if is_state:
         if mode.upper() == "PF":
@@ -142,7 +142,7 @@ def get_non_performing_loans_codes(states_or_region, mode="total"):
             non_performing_loans_by_location = NON_PERFORMING_LOANS_BY_REGION_PJ
 
     for location in states_or_region:
-        codes.append(non_performing_loans_by_location[location])
+        codes[location] = non_performing_loans_by_location[location]
     return codes
 
 

tests/sgs/test_regional_economy.py

@@ -5,22 +5,23 @@
 
 class TestGetNonPerformingLoansCodes:
     @pytest.mark.parametrize("states,expected_codes", [
-        (["ba", "pa"], ["15865", "15938"]),
-        (["BA"], ["15865"]),
-        ("N", ["15952"]),
+        (["ba", "pa"], {"BA": "15865", "PA": "15938"}),
+        (["BA"], {"BA": "15865"}),
+        ("N", {"N": "15952"}),
     ])
     def test_get_non_performing_loans_codes_by_state_total(self, states, expected_codes):
         assert get_non_performing_loans_codes(states) == expected_codes
 
 
 class TestGetNonPerformingLoans:
     @pytest.mark.parametrize("states,expected_columns", [
-        (["BA"], ["15865"]),
-        (["ba"], ["15865"]),
-        ("N", ["15888"]),
+        (["BA"], ["BA"]),
+        (["am", "pa"], ["AM", "PA"]),
+        ("N", ["N"]),
     ])
     def test_get_series_by_states_pf(self, states, expected_columns):
         series = get_non_performing_loans(states, last=10, mode="pf")
+
         assert isinstance(series, pd.DataFrame)
-        assert series.columns == expected_columns
+        assert (series.columns == expected_columns).all()
         assert len(series) == 10