1"""Utilities to manipulate JSON objects."""
2
3# NOTE: this is a copy of ipykernel/jsonutils.py (+blackified)
4
5# Copyright (c) IPython Development Team.
6# Distributed under the terms of the Modified BSD License.
7from __future__ import annotations
8
9import math
10import numbers
11import re
12import types
13from binascii import b2a_base64
14from datetime import datetime
15from typing import Any
16
17# -----------------------------------------------------------------------------
18# Globals and constants
19# -----------------------------------------------------------------------------
20
21# timestamp formats
22ISO8601 = "%Y-%m-%dT%H:%M:%S.%f"
23ISO8601_PAT = re.compile(
24 r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d{1,6})?Z?([\+\-]\d{2}:?\d{2})?$"
25)
26
27# holy crap, strptime is not threadsafe.
28# Calling it once at import seems to help.
29datetime.strptime("2000-01-01", "%Y-%m-%d")
30
31# -----------------------------------------------------------------------------
32# Classes and functions
33# -----------------------------------------------------------------------------
34
35
36# constants for identifying png/jpeg data
37PNG = b"\x89PNG\r\n\x1a\n"
38# front of PNG base64-encoded
39PNG64 = b"iVBORw0KG"
40JPEG = b"\xff\xd8"
41# front of JPEG base64-encoded
42JPEG64 = b"/9"
43# constants for identifying gif data
44GIF_64 = b"R0lGODdh"
45GIF89_64 = b"R0lGODlh"
46# front of PDF base64-encoded
47PDF64 = b"JVBER"
48
49
50def encode_images(format_dict: dict[str, str]) -> dict[str, str]:
51 """b64-encodes images in a displaypub format dict
52
53 Perhaps this should be handled in json_clean itself?
54
55 Parameters
56 ----------
57
58 format_dict : dict
59 A dictionary of display data keyed by mime-type
60
61 Returns
62 -------
63
64 format_dict : dict
65 A copy of the same dictionary,
66 but binary image data ('image/png', 'image/jpeg' or 'application/pdf')
67 is base64-encoded.
68
69 """
70 return format_dict
71
72
73def json_clean(obj: Any) -> Any:
74 """Clean an object to ensure it's safe to encode in JSON.
75
76 Atomic, immutable objects are returned unmodified. Sets and tuples are
77 converted to lists, lists are copied and dicts are also copied.
78
79 Note: dicts whose keys could cause collisions upon encoding (such as a dict
80 with both the number 1 and the string '1' as keys) will cause a ValueError
81 to be raised.
82
83 Parameters
84 ----------
85 obj : any python object
86
87 Returns
88 -------
89 out : object
90
91 A version of the input which will not cause an encoding error when
92 encoded as JSON. Note that this function does not *encode* its inputs,
93 it simply sanitizes it so that there will be no encoding errors later.
94
95 """
96 # types that are 'atomic' and ok in json as-is.
97 atomic_ok = (str, type(None))
98
99 # containers that we need to convert into lists
100 container_to_list = (tuple, set, types.GeneratorType)
101
102 # Since bools are a subtype of Integrals, which are a subtype of Reals,
103 # we have to check them in that order.
104
105 if isinstance(obj, bool):
106 return obj
107
108 if isinstance(obj, numbers.Integral):
109 # cast int to int, in case subclasses override __str__ (e.g. boost enum, #4598)
110 return int(obj)
111
112 if isinstance(obj, numbers.Real):
113 # cast out-of-range floats to their reprs
114 if math.isnan(obj) or math.isinf(obj):
115 return repr(obj)
116 return float(obj)
117
118 if isinstance(obj, atomic_ok):
119 return obj
120
121 if isinstance(obj, bytes):
122 return b2a_base64(obj).decode("ascii")
123
124 if isinstance(obj, container_to_list) or (
125 hasattr(obj, "__iter__") and hasattr(obj, "__next__")
126 ):
127 obj = list(obj)
128
129 if isinstance(obj, list):
130 return [json_clean(x) for x in obj]
131
132 if isinstance(obj, dict):
133 # First, validate that the dict won't lose data in conversion due to
134 # key collisions after stringification. This can happen with keys like
135 # True and 'true' or 1 and '1', which collide in JSON.
136 nkeys = len(obj)
137 nkeys_collapsed = len(set(map(str, obj)))
138 if nkeys != nkeys_collapsed:
139 raise ValueError(
140 "dict cannot be safely converted to JSON: "
141 "key collision would lead to dropped values"
142 )
143 # If all OK, proceed by making the new dict that will be json-safe
144 out = {}
145 for k, v in iter(obj.items()):
146 out[str(k)] = json_clean(v)
147 return out
148 if isinstance(obj, datetime):
149 return obj.strftime(ISO8601)
150
151 # we don't understand it, it's probably an unserializable object
152 raise ValueError("Can't clean for JSON: %r" % obj)