pumpwood_communication.serializers

  1"""Miscellaneous to help with serializers in communication."""
  2import base64
  3import simplejson as json
  4import orjson
  5import numpy as np
  6import pandas as pd
  7from typing import List, Union, Dict, Any
  8from simplejson import JSONEncoder
  9from datetime import datetime
 10from datetime import date
 11from datetime import time
 12from pandas import Timestamp
 13from shapely.geometry.base import BaseGeometry
 14from shapely.geometry import mapping
 15from sqlalchemy_utils.types.choice import Choice
 16from pumpwood_communication.exceptions import (
 17    PumpWoodException, PumpWoodNotImplementedError)
 18
 19
 20def default_encoder(obj):
 21    """Serialize complex objects."""
 22    # Return None if object is NaN
 23    if isinstance(obj, datetime):
 24        return obj.isoformat()
 25    if isinstance(obj, Timestamp):
 26        return obj.isoformat()
 27    if isinstance(obj, date):
 28        return obj.isoformat()
 29    if isinstance(obj, time):
 30        return obj.isoformat()
 31    if isinstance(obj, np.ndarray):
 32        return obj.tolist()
 33    if isinstance(obj, pd.DataFrame):
 34        return obj.to_dict('records')
 35    if isinstance(obj, pd.Series):
 36        obj.dtype == Timestamp
 37        return obj.tolist()
 38    if isinstance(obj, np.generic):
 39        return obj.item()
 40    if isinstance(obj, BaseGeometry):
 41        if obj.is_empty:
 42            return None
 43        else:
 44            return mapping(obj)
 45    if isinstance(obj, BaseGeometry):
 46        return mapping(obj)
 47    if isinstance(obj, Choice):
 48        return obj.code
 49    if isinstance(obj, set):
 50        return list(obj)
 51    else:
 52        raise TypeError(
 53            "Unserializable object {} of type {}".format(obj, type(obj)))
 54
 55
 56class PumpWoodJSONEncoder(JSONEncoder):
 57    """PumpWood default serializer.
 58
 59    Treat not simple python types to facilitate at serialization of
 60    pandas, numpy, data, datetime and other data types.
 61    """
 62
 63    def default(self, obj):
 64        """Serialize complex objects."""
 65        return default_encoder(obj)
 66
 67
 68def pumpJsonDump(x: any, sort_keys: bool = False,  # NOQA
 69                 indent: Union[int, bool] = None):
 70    """Dump a Json to python object.
 71
 72    Args:
 73        x (any):
 74            Object to be serialized using PumpWoodJSONEncoder encoder.
 75        sort_keys (bool):
 76            If json serialized data should have its keys sorted. This option
 77            makes serialization return of data reproductable.
 78        indent (int):
 79            Pass indent argument to simplejson dumps.
 80    """
 81    # Compatibility with simplejson serialization
 82    is_indent = indent is not None
 83    if sort_keys and is_indent:
 84        return orjson.dumps(x, default=default_encoder, option=(
 85            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
 86            orjson.OPT_SORT_KEYS | orjson.OPT_INDENT_2 |
 87            orjson.OPT_SERIALIZE_NUMPY))
 88    elif sort_keys:
 89        return orjson.dumps(x, default=default_encoder, option=(
 90            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
 91            orjson.OPT_SORT_KEYS | orjson.OPT_SERIALIZE_NUMPY))
 92    elif is_indent:
 93        return orjson.dumps(x, default=default_encoder, option=(
 94            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
 95            orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY))
 96    else:
 97        return orjson.dumps(x, default=default_encoder, option=(
 98            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
 99            orjson.OPT_SERIALIZE_NUMPY))
100
101
102class CompositePkBase64Converter:
103    """Convert composite primary keys in base64 dictionary."""
104
105    @staticmethod
106    def get_attribute(obj: Any, att: str) -> Any:
107        """Get attribute from object or dictinary.
108
109        Args:
110            obj (Any):
111                Object or a dictinary.
112            att (str):
113                Name of the attribute/key that will be used to return
114                the value.
115
116        Return:
117            Return object/dictionary value associated with attribute.
118        """
119        temp_pk_value = None
120        if type(obj) is dict:
121            temp_pk_value = obj[att]
122        else:
123            temp_pk_value = getattr(obj, att)
124        return temp_pk_value
125
126    @classmethod
127    def dump(cls, obj, primary_keys: Union[str, List[str], Dict[str, str]]
128             ) -> Union[str, int]:
129        """Convert primary keys and composite to a single value.
130
131        Treat cases when more than one column are used as primary keys,
132        at this cases, a base64 used on url serialization of the dictionary
133        is returned.
134
135        Args:
136            obj:
137                SQLAlchemy object.
138            primary_keys (Union[str, List[str], Dict[str, str]):
139                As string, a list or a dictionary leading to different
140                behaviour.
141                - **str:** It will return the value associated with object
142                    attribute.
143                - **List[str]:** If list has lenght equal to 1, it will have
144                    same behaviour as str. If greater than 1, it will be
145                    returned a base64 encoded dictionary with the keys at
146                    primary_keys.
147                - **Dict[str, str]:** Dictionary to map object fields to
148                    other keys. This is usefull when querying related fields
149                    by composite forenging keys to match original data fieds.
150
151        Returns:
152            If the primary key is unique, return the value of the primary
153            key, if is have more than one column as primary key, return
154            a dictionary of the primary keys encoded as base64 url safe.
155        """
156        if type(primary_keys) is str:
157            return getattr(obj, primary_keys)
158
159        elif type(primary_keys) is list:
160            if len(primary_keys) == 1:
161                return getattr(obj, primary_keys[0])
162            else:
163                # Will return a None value if all composite primary keys are
164                # None
165                is_all_none = False
166                composite_pk_dict = {}
167                for pk_col in primary_keys:
168                    temp_pk_value = cls.get_attribute(obj, pk_col)
169                    is_all_none = is_all_none or (temp_pk_value is None)
170                    composite_pk_dict[pk_col] = temp_pk_value
171                if is_all_none:
172                    return None
173
174                # If not all primary keys are None, them serialize it and
175                # convert to a base64 dictionary to be used as PK
176                composite_pk_str = pumpJsonDump(composite_pk_dict)
177                return base64.urlsafe_b64encode(
178                    composite_pk_str.encode()).decode()
179
180        # Map object values to other, this is used when builds forenging
181        # key references and request related field using microservice.
182        elif type(primary_keys) is dict:
183            # Will return a None value if all composite primary keys are
184            # None
185            is_all_none = False
186            composite_pk_dict = {}
187            for key, value in primary_keys.items():
188                temp_pk_value = cls.get_attribute(obj, key)
189                is_all_none = is_all_none or (temp_pk_value is None)
190                composite_pk_dict[value] = temp_pk_value
191            if is_all_none:
192                return None
193
194            # If not all primary keys are None, them serialize it and
195            # convert to a base64 dictionary to be used as PK. Using
196            # dictionary will map values before converting to base64
197            # dictionary
198            composite_pk_str = pumpJsonDump(composite_pk_dict)
199            base64_composite_pk = base64.urlsafe_b64encode(
200                composite_pk_str.encode()).decode()
201            return base64_composite_pk
202
203        # This will raise error if primary_keys type is not implemented
204        else:
205            msg = (
206                "CompositePkBase64Converter.dump argument primary_keys "
207                "is not a list of strings or a map dictionary. Type "
208                "[{arg_type}]").format(arg_type=type(primary_keys))
209            raise PumpWoodNotImplementedError(message=msg)
210
211    @staticmethod
212    def load(value: Union[str, int]) -> Union[int, dict]:
213        """Convert encoded primary keys to values.
214
215        If the primary key is a string, try to transform it to dictionary
216        decoding json base64 to a dictionary.
217
218        Args:
219            value:
220                Primary key value as an integer or as a base64
221                encoded json dictionary.
222
223        Return:
224            Return the primary key as integer if possible, or try to decoded
225            it to a dictionary from a base64 encoded json.
226        """
227        # Try to convert value to integer
228        try:
229            float_value = float(value)
230            if float_value.is_integer():
231                return int(float_value)
232            else:
233                msg = "[{value}] value is a float, but not integer."
234                raise PumpWoodException(msg, payload={"value": value})
235
236        # If not possible, try to decode a base64 JSON dictionary
237        except Exception as e1:
238            try:
239                return json.loads(base64.b64decode(value))
240            except Exception as e2:
241                msg = (
242                    "[{value}] value is not an integer and could no be "
243                    "decoded as a base64 encoded json dictionary. Value=")
244                raise PumpWoodException(
245                    message=msg, payload={
246                        "value": value,
247                        "exception_int": str(e1),
248                        "exception_base64": str(e2)})