pumpwood_communication.misc

 1"""Miscelaneus function to help in development."""
 2import pandas as pd
 3import copy
 4from typing import List, Literal
 5from pumpwood_communication.exceptions import PumpWoodException
 6
 7
 8def unpack_dict_columns(df: pd.DataFrame, columns: List[str]) -> pd.DataFrame:
 9    """Unpack dictinary columns at a dataframe.
10
11    Return a copy of the dataframe with 'columns' unpacked and removed
12    from result.
13
14    Args:
15        df (pd.DataFrame):
16            Dataframe to unpack the columns.
17        columns (List[str]):
18            List of columns to unpack in columns.
19
20    Return:
21        Return a dataframe with dict columns unpacked.
22    """
23    list_unpacked_results = []
24    for c in columns:
25        list_unpacked_results.append(df[c].apply(pd.Series))
26    return pd.concat(
27        [df.drop(columns=columns)] + list_unpacked_results,
28        axis=1)
29
30
31def extract_dict_subset(data: dict, keys: List[str],
32                        raise_not_present: Literal[
33                            'raise', 'ignore', 'add_none'] = 'raise'):
34    """Extract a subset of data from dictonary.
35
36    Args:
37        data (dict):
38            Dictionary data from which the subset will be extracted.
39        keys (List[str]):
40            Keys that will be extracted.
41        raise_not_present (str):
42            Control if an erros should be raised when key is not present.
43            - 'raise': Raise error
44            - 'ignore': Does not return the missing key on function result.
45            - 'add_none': Return key with None value.
46
47    Returns:
48        Return a dictonary with a copy of subset of the keys of original
49        dictonary.
50    """
51    if raise_not_present not in ['raise', 'ignore', 'add_none']:
52        msg = (
53            'raise_not_present must be in [raise, ignore, add_none].'
54            'raise_not_present={raise_not_present}')
55        raise PumpWoodException(
56            message=msg, payload={
57                'raise_not_present': raise_not_present})
58
59    return_dict = {}
60    for key in keys:
61        temp_data = data.get(key)
62        if temp_data is None:
63            if raise_not_present == 'raise':
64                msg = (
65                    'key [{key}] not found on dictonary and raise_not_present '
66                    'arg is set as [{raise_not_present}]')
67                raise PumpWoodException(
68                    message=msg, payload={
69                        'key': key, 'raise_not_present': raise_not_present})
70            if raise_not_present == 'ignore':
71                continue
72            if raise_not_present == 'add_none':
73                return_dict[key] = None
74        else:
75            return_dict[key] = copy.deepcopy(temp_data)
76    return return_dict