1  import datetime
 2  import json
 3  import typing
 4  import uuid
 5  from functools import partial
 6  from typing import Callable
 7  from typing import Tuple
 8  from typing import Type
 9  
10  import numpy as np
11  import pandas as pd
12  
13  from evidently.legacy.core import ColumnType
14  from evidently.legacy.utils.types import ApproxValue
15  
16  _TYPES_MAPPING = (
17      (
18          (np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64),
19          int,
20      ),
21      ((np.double, np.float16, np.float32, np.float64), lambda obj: None if obj is np.nan else float(obj)),
22      ((np.ndarray,), lambda obj: obj.tolist()),
23      ((np.bool_), bool),
24      ((pd.Timedelta,), str),
25      ((np.void, type(pd.NaT)), lambda obj: None),  # should be before datetime as NaT is subclass of datetime.
26      ((pd.Timestamp, datetime.datetime, datetime.date), lambda obj: obj.isoformat()),
27      # map ApproxValue to json value
28      ((ApproxValue,), lambda obj: obj.dict()),
29      ((pd.Series, pd.Index, pd.Categorical), lambda obj: obj.tolist()),
30      ((pd.DataFrame,), lambda obj: obj.to_dict()),
31      ((frozenset,), lambda obj: list(obj)),
32      ((uuid.UUID,), lambda obj: str(obj)),
33      ((ColumnType,), lambda obj: obj.value),
34      ((pd.Period,), lambda obj: str(obj)),
35  )
36  
37  
38  def add_type_mapping(types: Tuple[Type], encoder: Callable):
39      global _TYPES_MAPPING  # noqa: PLW0603
40      _TYPES_MAPPING += ((types, encoder),)  # type: ignore[assignment]
41  
42  
43  class NumpyEncoder(json.JSONEncoder):
44      """Numpy and Pandas data types to JSON types encoder"""
45  
46      def default(self, o):
47          """JSON converter calls the method when it cannot convert an object to a Python type
48          Convert the object to a Python type
49  
50          If we cannot convert the object, leave the default `JSONEncoder` behaviour - raise a TypeError exception.
51          """
52  
53          # check mapping rules
54          for types_list, python_type in _TYPES_MAPPING:
55              if isinstance(o, types_list):
56                  return python_type(o)
57  
58          # explicit check pandas null
59          if not isinstance(o, typing.Sequence) and pd.isnull(o):
60              return None
61  
62          return json.JSONEncoder.default(self, o)
63  
64  
65  numpy_dumps = partial(json.dumps, cls=NumpyEncoder)