1"""Utilities to manipulate JSON objects."""
2
3# Copyright (c) Jupyter Development Team.
4# Distributed under the terms of the Modified BSD License.
5
6from datetime import datetime
7import re
8import warnings
9
10from dateutil.parser import parse as _dateutil_parse
11from dateutil.tz import tzlocal
12
13next_attr_name = '__next__' # Not sure what downstream library uses this, but left it to be safe
14
15#-----------------------------------------------------------------------------
16# Globals and constants
17#-----------------------------------------------------------------------------
18
19# timestamp formats
20ISO8601 = "%Y-%m-%dT%H:%M:%S.%f"
21ISO8601_PAT = re.compile(r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d{1,6})?(Z|([\+\-]\d{2}:?\d{2}))?$")
22
23# holy crap, strptime is not threadsafe.
24# Calling it once at import seems to help.
25datetime.strptime("1", "%d")
26
27#-----------------------------------------------------------------------------
28# Classes and functions
29#-----------------------------------------------------------------------------
30
31def _ensure_tzinfo(dt):
32    """Ensure a datetime object has tzinfo
33
34    If no tzinfo is present, add tzlocal
35    """
36    if not dt.tzinfo:
37        # No more naïve datetime objects!
38        warnings.warn("Interpreting naive datetime as local %s. Please add timezone info to timestamps." % dt,
39            DeprecationWarning,
40            stacklevel=4)
41        dt = dt.replace(tzinfo=tzlocal())
42    return dt
43
44def parse_date(s):
45    """parse an ISO8601 date string
46
47    If it is None or not a valid ISO8601 timestamp,
48    it will be returned unmodified.
49    Otherwise, it will return a datetime object.
50    """
51    if s is None:
52        return s
53    m = ISO8601_PAT.match(s)
54    if m:
55        dt = _dateutil_parse(s)
56        return _ensure_tzinfo(dt)
57    return s
58
59def extract_dates(obj):
60    """extract ISO8601 dates from unpacked JSON"""
61    if isinstance(obj, dict):
62        new_obj = {} # don't clobber
63        for k,v in obj.items():
64            new_obj[k] = extract_dates(v)
65        obj = new_obj
66    elif isinstance(obj, (list, tuple)):
67        obj = [ extract_dates(o) for o in obj ]
68    elif isinstance(obj, str):
69        obj = parse_date(obj)
70    return obj
71
72def squash_dates(obj):
73    """squash datetime objects into ISO8601 strings"""
74    if isinstance(obj, dict):
75        obj = dict(obj) # don't clobber
76        for k,v in obj.items():
77            obj[k] = squash_dates(v)
78    elif isinstance(obj, (list, tuple)):
79        obj = [ squash_dates(o) for o in obj ]
80    elif isinstance(obj, datetime):
81        obj = obj.isoformat()
82    return obj
83
84def date_default(obj):
85    """default function for packing datetime objects in JSON."""
86    if isinstance(obj, datetime):
87        obj = _ensure_tzinfo(obj)
88        return obj.isoformat().replace('+00:00', 'Z')
89    else:
90        raise TypeError("%r is not JSON serializable" % obj)
91
92