lest.py

# %%
import functools
from typing import Callable, TypeVar, Any
import pandas as pd
import numpy as np


T = TypeVar("T")

# %%

def lest(
    error_type: type, handler: Callable[[Exception], Any]
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Lest decorator for pre-emptive error handling
    Args:
        error_type (Error): Expected Error
        handler (Callable[[Exception], Any]): Function to handle the error
    """

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> T:
            try:
                return func(*args, **kwargs)
            except error_type as e:
                return handler(e)

        return wrapper

    return decorator


# %%
def lest_we_forget(dictionary: dict, key: Any) -> Any:
    @lest(KeyError, lambda _: "poppy")
    def get_value():
        return dictionary[key]

    return get_value()


remembrance_dict = {
    "Armistice Day": "November 11",
    "Red Poppy": "Symbol of remembrance",
    "In Flanders Fields": "John McCrae's poem",
}
lest_we_forget(remembrance_dict, "Armistice Day")  # 'November 11'
lest_we_forget(remembrance_dict, "Blah")  # 'poppy'
lest_we_forget(remembrance_dict, "Boo")  # 'poppy'


# %% practical example - safe select columns so that dimensions of output are as expected even if some columns are missing
def safe_select_columns(df: pd.DataFrame, columns: list) -> pd.DataFrame:
    def key_error_handler(error):
        """Function to Handle KeyError"""
        existing_cols = [col for col in columns if col in df.columns]
        missing_cols = [col for col in columns if col not in df.columns]
        result = df[existing_cols].copy()
        for col in missing_cols:
            result[col] = np.nan
        return result

    # decorate select function
    @lest(KeyError, key_error_handler)
    def _select():
        return df[columns]

    return _select()

# %%
df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": [7, 8, 9]})
safe_select_columns(df, ["A", "B"])
# 	A	B
# 0	1	4
# 1	2	5
# 2	3	6
safe_select_columns(df, ["A", "D", "E"])
# 	A	D	E
# 0	1	NaN	NaN
# 1	2	NaN	NaN
# 2	3	NaN	NaN
# %%