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 decimal import Decimal
 13from pandas import Timestamp
 14from shapely.geometry.base import BaseGeometry
 15from shapely.geometry import mapping
 16from sqlalchemy_utils.types.choice import Choice
 17from pumpwood_communication.exceptions import (
 18    PumpWoodException, PumpWoodNotImplementedError)
 19from pumpwood_communication.type import (
 20    PumpwoodSentinel, PumpwoodDataclassMixin, MISSING)
 21
 22
 23def default_encoder(obj):
 24    """Serialize complex objects."""
 25    # Return None if object is NaN
 26    if not isinstance(obj, (pd.DataFrame, pd.Series, np.ndarray, list, dict)):
 27        if pd.isna(obj):
 28            return None
 29
 30    if isinstance(obj, (datetime, Timestamp, date, time)):
 31        return obj.isoformat()
 32    if isinstance(obj, np.ndarray):
 33        return obj.tolist()
 34    if isinstance(obj, pd.DataFrame):
 35        return obj.to_dict('records')
 36    if isinstance(obj, pd.Series):
 37        return obj.tolist()
 38    if isinstance(obj, np.generic):
 39        return obj.item()
 40    if isinstance(obj, Decimal):
 41        return float(obj)
 42    if isinstance(obj, BaseGeometry):
 43        if obj.is_empty:
 44            return None
 45        else:
 46            return mapping(obj)
 47    if isinstance(obj, Choice):
 48        return obj.code
 49    if isinstance(obj, set):
 50        return list(obj)
 51
 52    #########################################################
 53    # TODO: Adjust convertion of decimal to preseve precision
 54    # There is lost of precision when converting decimal to float,
 55    # but Decimal is not currently parsiable using orjson
 56    if isinstance(obj, Decimal):
 57        return float(obj)
 58
 59    ###################################
 60    # Serialize Pumpwood expecial types
 61    if isinstance(obj, PumpwoodDataclassMixin):
 62        return obj.to_dict()
 63    if isinstance(obj, PumpwoodSentinel):
 64        return obj.value()
 65    else:
 66        raise TypeError(
 67            "Unserializable object {} of type {}".format(obj, type(obj)))
 68
 69
 70class PumpWoodJSONEncoder(JSONEncoder):
 71    """PumpWood default serializer.
 72
 73    Treat not simple python types to facilitate at serialization of
 74    pandas, numpy, data, datetime and other data types.
 75    """
 76
 77    def default(self, obj):
 78        """Serialize complex objects."""
 79        return default_encoder(obj)
 80
 81
 82def pumpJsonDump(x: any, sort_keys: bool = False,  # NOQA
 83                 indent: Union[int, bool] = None):
 84    """Dump a Json to python object.
 85
 86    Args:
 87        x (any):
 88            Object to be serialized using PumpWoodJSONEncoder encoder.
 89        sort_keys (bool):
 90            If json serialized data should have its keys sorted. This option
 91            makes serialization return of data reproductable.
 92        indent (int):
 93            Pass indent argument to simplejson dumps.
 94    """
 95    # Compatibility with simplejson serialization
 96    is_indent = indent is not None
 97    if sort_keys and is_indent:
 98        return orjson.dumps(x, default=default_encoder, option=(
 99            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
100            orjson.OPT_SORT_KEYS | orjson.OPT_INDENT_2 |
101            orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_PASSTHROUGH_DATACLASS))
102    elif sort_keys:
103        return orjson.dumps(x, default=default_encoder, option=(
104            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
105            orjson.OPT_SORT_KEYS | orjson.OPT_SERIALIZE_NUMPY |
106            orjson.OPT_PASSTHROUGH_DATACLASS))
107    elif is_indent:
108        return orjson.dumps(x, default=default_encoder, option=(
109            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
110            orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY |
111            orjson.OPT_PASSTHROUGH_DATACLASS))
112    else:
113        return orjson.dumps(x, default=default_encoder, option=(
114            orjson.OPT_NAIVE_UTC | orjson.OPT_NON_STR_KEYS |
115            orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_PASSTHROUGH_DATACLASS))
116
117
118class CompositePkBase64Converter:
119    """Convert composite primary keys in base64 dictionary."""
120
121    @staticmethod
122    def get_attribute(obj: Any, att: str) -> Any:
123        """Get attribute from object or dictinary.
124
125        Args:
126            obj (Any):
127                Object or a dictinary.
128            att (str):
129                Name of the attribute/key that will be used to return
130                the value.
131
132        Return:
133            Return object/dictionary value associated with attribute.
134        """
135        if type(obj) is dict:
136            return obj.get(att, MISSING)
137        else:
138            return getattr(obj, att, MISSING)
139
140    @classmethod
141    def dump(cls, obj: object | dict,
142             primary_keys: Union[str, List[str], Dict[str, str]]
143             ) -> Union[str, int]:
144        """Convert primary keys and composite to a single value.
145
146        Treat cases when more than one column are used as primary keys,
147        at this cases, a base64 used on url serialization of the dictionary
148        is returned.
149
150        Args:
151            obj:
152                SQLAlchemy object or dictionary with data to build the
153                forenging key.
154            primary_keys (Union[str, List[str], Dict[str, str]):
155                As string, a list or a dictionary leading to different
156                behaviour.
157                - **str:** It will return the value associated with object
158                    attribute.
159                - **List[str]:** If list has length equal to 1, it will have
160                    same behaviour as str. If greater than 1, it will be
161                    returned a base64 encoded dictionary with the keys at
162                    primary_keys.
163                - **Dict[str, str]:** Dictionary to map object fields to
164                    other keys. This is usefull when querying related fields
165                    by composite forenging keys to match original data fieds.
166
167        Returns:
168            If the primary key is unique, return the value of the primary
169            key, if is have more than one column as primary key, return
170            a dictionary of the primary keys encoded as base64 url safe.
171        """
172        if obj is None:
173            return None
174
175        missing_keys = []
176        return_value = None
177        if isinstance(primary_keys, str):
178            key_value = cls.get_attribute(obj, primary_keys)
179            if key_value != MISSING:
180                return_value = key_value
181            else:
182                missing_keys.append(primary_keys)
183
184        elif isinstance(primary_keys, list):
185            if len(primary_keys) == 1:
186                key_value = cls.get_attribute(obj, primary_keys[0])
187                if key_value != MISSING:
188                    return key_value
189                else:
190                    missing_keys.append(primary_keys[0])
191            else:
192                composite_pk_dict = {}
193                for pk_col in primary_keys:
194                    key_value = cls.get_attribute(obj, pk_col)
195                    if key_value != MISSING:
196                        composite_pk_dict[pk_col] = key_value
197                    else:
198                        missing_keys.append(pk_col)
199                return_value = composite_pk_dict
200
201        # Map object values to other, this is used when builds forenging
202        # key references and request related field using microservice.
203        elif isinstance(primary_keys, dict):
204            # Treat the case when the dictinary is only an id->value
205            composite_pk_dict = {}
206            for key, value in primary_keys.items():
207                key_value = cls.get_attribute(obj, key)
208                if key_value != MISSING:
209                    composite_pk_dict[value] = key_value
210                else:
211                    missing_keys.append(key)
212            return_value = composite_pk_dict
213
214        # Check if some missing keys are there
215        if len(missing_keys) != 0:
216            msg = (
217                "Some keys were not found on object/dict to create "
218                "the composite forenging key. "
219                "Primary Keys: {primary_keys}; "
220                "Missing keys: {missing_keys}")
221            raise PumpWoodException(msg, payload={
222                "primary_keys": list(primary_keys.keys()),
223                "missing_keys": missing_keys})
224
225        if isinstance(return_value, dict):
226            if {'id'} == set(return_value.keys()):
227                return return_value['id']            
228            return cls.dump_dict(primary_key_dict=return_value)
229        else:
230            return return_value
231
232    @classmethod
233    def validate_primary_key_dict(cls, primary_key_dict: dict):
234        """."""
235        if not isinstance(primary_key_dict, dict):
236            msg = "primary_key_dict must be a dictionary. Received: {type}"
237            raise PumpWoodNotImplementedError(
238                msg, payload={'type': type(primary_key_dict).__name__})
239
240        for key, value in primary_key_dict.items():
241            if isinstance(value, (dict, list, tuple)):
242                msg = (
243                    "primary_key_dict must be flat dictionary. "
244                    "Nested type[{type}] found at key [{key}]")
245                raise PumpWoodException(
246                    msg.format(),
247                    payload={
248                        'primary_key_dict': primary_key_dict,
249                        'type': type(value).__name__,
250                        'key': key})
251
252    @classmethod
253    def dump_dict(cls, primary_key_dict: dict) -> str:
254        """Dump a primary key dictionary to base64 url safe string.
255
256        Args:
257            primary_key_dict (dict):
258                Dictionary with the primary key values.
259
260        Returns:
261            str:
262                Base64 url safe string of the primary key dictionary.
263        """
264        cls.validate_primary_key_dict(primary_key_dict=primary_key_dict)
265        composite_pk_str = pumpJsonDump(primary_key_dict)
266        base64_composite_pk = base64.urlsafe_b64encode(composite_pk_str)\
267            .decode()
268        return base64_composite_pk
269
270    @staticmethod
271    def load(value: Union[str, int]) -> Union[int, dict]:
272        """Convert encoded primary keys to values.
273
274        If the primary key is a string, try to transform it to dictionary
275        decoding json base64 to a dictionary.
276
277        Args:
278            value:
279                Primary key value as an integer or as a base64
280                encoded json dictionary.
281
282        Return:
283            Return the primary key as integer if possible, or try to decoded
284            it to a dictionary from a base64 encoded json.
285        """
286        # Try to convert value to integer
287        try:
288            float_value = float(value)
289            if float_value.is_integer():
290                return int(float_value)
291            else:
292                msg = "[{value}] value is a float, but not integer."
293                raise PumpWoodException(msg, payload={"value": value})
294
295        # If not possible, try to decode a base64 JSON dictionary
296        except Exception as e1:
297            try:
298                return json.loads(base64.urlsafe_b64decode(value))
299            except Exception as e2:
300                msg = (
301                    "[{value}] value is not an integer and could no be "
302                    "decoded as a base64 encoded json dictionary. Value=")
303                raise PumpWoodException(
304                    message=msg, payload={
305                        "value": value,
306                        "exception_int": str(e1),
307                        "exception_base64": str(e2)})