From 2dbd0a62e80c964e0a279125d7cb06760d7272cd Mon Sep 17 00:00:00 2001
From: "j.aschauer" <j.aschauer@oeko.de>
Date: Thu, 23 Nov 2023 15:20:33 +0100
Subject: [PATCH] refactor select subregions to function

---
 app/plot_functions.py   |  8 ++++++--
 app/ptxboa_functions.py | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/app/plot_functions.py b/app/plot_functions.py
index 274bede2..207ed705 100644
--- a/app/plot_functions.py
+++ b/app/plot_functions.py
@@ -9,7 +9,11 @@
 import plotly.graph_objects as go
 import streamlit as st
 
-from app.ptxboa_functions import remove_subregions, subset_and_pivot_input_data
+from app.ptxboa_functions import (
+    remove_subregions,
+    select_subregions,
+    subset_and_pivot_input_data,
+)
 from ptxboa.api import PtxboaAPI
 
 
@@ -226,7 +230,7 @@ def _choropleth_map_deep_dive_country(
     if custom_data_func_kwargs is None:
         custom_data_func_kwargs = {}
     # subsetting 'df' for the selected deep dive country
-    df = df.copy().loc[df.index.str.startswith(f"{deep_dive_country} ("), :]
+    df = select_subregions(df, deep_dive_country)
     # need to calculate custom data befor is03166 column is appended.
     hover_data = custom_data_func(df, **custom_data_func_kwargs)
     # get dataframe with info about iso 3166-2 codes and map them to res_costs
diff --git a/app/ptxboa_functions.py b/app/ptxboa_functions.py
index ca9f3cdb..d9fae213 100644
--- a/app/ptxboa_functions.py
+++ b/app/ptxboa_functions.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """Utility functions for streamlit app."""
+from typing import Literal
 
 import pandas as pd
 import streamlit as st
@@ -190,6 +191,27 @@ def remove_subregions(api: PtxboaAPI, df: pd.DataFrame, country_name: str):
     return df
 
 
+def select_subregions(
+    df: pd.DataFrame, deep_dive_country: Literal["Argentina", "Morocco", "South Africa"]
+) -> pd.DataFrame:
+    """
+    Only select rows corresponding to subregions of a deep dive country.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        pandas DataFrame with list of regions as index.
+    deep_dive_country : str in {"Argentina", "Morocco", "South Africa"}
+
+
+    Returns
+    -------
+    pd.DataFrame
+    """
+    df = df.copy().loc[df.index.str.startswith(f"{deep_dive_country} ("), :]
+    return df
+
+
 def reset_user_changes():
     """Reset all user changes."""
     if (