jsonutil.py 2.9 KB

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