pumpwood_communication.misc

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