1"""Utilities to manipulate JSON objects."""
2
3# Copyright (c) Jupyter Development Team.
4# Distributed under the terms of the Modified BSD License.
5import math
6import numbers
7import re
8import types
9import warnings
10from binascii import b2a_base64
11from collections.abc import Iterable
12from datetime import date, datetime
13from typing import Any, Union
14
15from dateutil.parser import isoparse as _dateutil_parse
16from dateutil.tz import tzlocal
17
18next_attr_name = "__next__" # Not sure what downstream library uses this, but left it to be safe
19
20# -----------------------------------------------------------------------------
21# Globals and constants
22# -----------------------------------------------------------------------------
23
24# timestamp formats
25ISO8601 = "%Y-%m-%dT%H:%M:%S.%f"
26ISO8601_PAT = re.compile(
27 r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d{1,6})?(Z|([\+\-]\d{2}:?\d{2}))?$"
28)
29
30# holy crap, strptime is not threadsafe.
31# Calling it once at import seems to help.
32datetime.strptime("2000-01-01", "%Y-%m-%d") # noqa
33
34# -----------------------------------------------------------------------------
35# Classes and functions
36# -----------------------------------------------------------------------------
37
38
39def _ensure_tzinfo(dt: datetime) -> datetime:
40 """Ensure a datetime object has tzinfo
41
42 If no tzinfo is present, add tzlocal
43 """
44 if not dt.tzinfo:
45 # No more naïve datetime objects!
46 warnings.warn(
47 "Interpreting naive datetime as local %s. Please add timezone info to timestamps." % dt,
48 DeprecationWarning,
49 stacklevel=4,
50 )
51 dt = dt.replace(tzinfo=tzlocal())
52 return dt
53
54
55def parse_date(s: str | None) -> Union[str, datetime] | None:
56 """parse an ISO8601 date string
57
58 If it is None or not a valid ISO8601 timestamp,
59 it will be returned unmodified.
60 Otherwise, it will return a datetime object.
61 """
62 if s is None:
63 return s
64 m = ISO8601_PAT.match(s)
65 if m:
66 dt = _dateutil_parse(s)
67 return _ensure_tzinfo(dt)
68 return s
69
70
71def extract_dates(obj: Any) -> Any:
72 """extract ISO8601 dates from unpacked JSON"""
73 if isinstance(obj, dict):
74 new_obj = {} # don't clobber
75 for k, v in obj.items():
76 new_obj[k] = extract_dates(v)
77 obj = new_obj
78 elif isinstance(obj, list | tuple):
79 obj = [extract_dates(o) for o in obj]
80 elif isinstance(obj, str):
81 obj = parse_date(obj)
82 return obj
83
84
85def squash_dates(obj: Any) -> Any:
86 """squash datetime objects into ISO8601 strings"""
87 if isinstance(obj, dict):
88 obj = dict(obj) # don't clobber
89 for k, v in obj.items():
90 obj[k] = squash_dates(v)
91 elif isinstance(obj, list | tuple):
92 obj = [squash_dates(o) for o in obj]
93 elif isinstance(obj, datetime):
94 obj = obj.isoformat()
95 return obj
96
97
98def date_default(obj: Any) -> Any:
99 """DEPRECATED: Use jupyter_client.jsonutil.json_default"""
100 warnings.warn(
101 "date_default is deprecated since jupyter_client 7.0.0."
102 " Use jupyter_client.jsonutil.json_default.",
103 stacklevel=2,
104 )
105 return json_default(obj)
106
107
108def json_default(obj: Any) -> Any:
109 """default function for packing objects in JSON."""
110 if isinstance(obj, datetime):
111 obj = _ensure_tzinfo(obj)
112 return obj.isoformat().replace("+00:00", "Z")
113
114 if isinstance(obj, date):
115 return obj.isoformat()
116
117 if isinstance(obj, bytes):
118 return b2a_base64(obj, newline=False).decode("ascii")
119
120 if isinstance(obj, Iterable):
121 return list(obj)
122
123 if isinstance(obj, numbers.Integral):
124 return int(obj)
125
126 if isinstance(obj, numbers.Real):
127 return float(obj)
128
129 raise TypeError("%r is not JSON serializable" % obj)
130
131
132# Copy of the old ipykernel's json_clean
133# This is temporary, it should be removed when we deprecate support for
134# non-valid JSON messages
135def json_clean(obj: Any) -> Any:
136 # types that are 'atomic' and ok in json as-is.
137 atomic_ok = (str, type(None))
138
139 # containers that we need to convert into lists
140 container_to_list = (tuple, set, types.GeneratorType)
141
142 # Since bools are a subtype of Integrals, which are a subtype of Reals,
143 # we have to check them in that order.
144
145 if isinstance(obj, bool):
146 return obj
147
148 if isinstance(obj, numbers.Integral):
149 # cast int to int, in case subclasses override __str__ (e.g. boost enum, #4598)
150 return int(obj)
151
152 if isinstance(obj, numbers.Real):
153 # cast out-of-range floats to their reprs
154 if math.isnan(obj) or math.isinf(obj):
155 return repr(obj)
156 return float(obj)
157
158 if isinstance(obj, atomic_ok):
159 return obj
160
161 if isinstance(obj, bytes):
162 # unanmbiguous binary data is base64-encoded
163 # (this probably should have happened upstream)
164 return b2a_base64(obj, newline=False).decode("ascii")
165
166 if isinstance(obj, container_to_list) or (
167 hasattr(obj, "__iter__") and hasattr(obj, next_attr_name)
168 ):
169 obj = list(obj)
170
171 if isinstance(obj, list):
172 return [json_clean(x) for x in obj]
173
174 if isinstance(obj, dict):
175 # First, validate that the dict won't lose data in conversion due to
176 # key collisions after stringification. This can happen with keys like
177 # True and 'true' or 1 and '1', which collide in JSON.
178 nkeys = len(obj)
179 nkeys_collapsed = len(set(map(str, obj)))
180 if nkeys != nkeys_collapsed:
181 msg = (
182 "dict cannot be safely converted to JSON: "
183 "key collision would lead to dropped values"
184 )
185 raise ValueError(msg)
186 # If all OK, proceed by making the new dict that will be json-safe
187 out = {}
188 for k, v in obj.items():
189 out[str(k)] = json_clean(v)
190 return out
191
192 if isinstance(obj, datetime | date):
193 return obj.strftime(ISO8601)
194
195 # we don't understand it, it's probably an unserializable object
196 raise ValueError("Can't clean for JSON: %r" % obj)