1#!/usr/bin/env python
2"""Json related utilities."""
3import copy
4import datetime
5import logging
6
7try:
8  import json
9except ImportError:
10  import simplejson as json
11
12from google.appengine.api import datastore_errors
13from google.appengine.api import datastore_types
14from google.appengine.ext import db
15from google.appengine.ext import ndb
16
17# pylint: disable=invalid-name
18
19
20class JsonEncoder(json.JSONEncoder):
21  """MR customized json encoder."""
22
23  TYPE_ID = "__mr_json_type"
24
25  def default(self, o):
26    """Inherit docs."""
27    if type(o) in _TYPE_TO_ENCODER:
28      encoder = _TYPE_TO_ENCODER[type(o)]
29      json_struct = encoder(o)
30      json_struct[self.TYPE_ID] = type(o).__name__
31      return json_struct
32    return super(JsonEncoder, self).default(o)
33
34
35class JsonDecoder(json.JSONDecoder):
36  """MR customized json decoder."""
37
38  def __init__(self, **kwargs):
39    if "object_hook" not in kwargs:
40      kwargs["object_hook"] = self._dict_to_obj
41    super(JsonDecoder, self).__init__(**kwargs)
42
43  def _dict_to_obj(self, d):
44    """Converts a dictionary of json object to a Python object."""
45    if JsonEncoder.TYPE_ID not in d:
46      return d
47
48    type_name = d.pop(JsonEncoder.TYPE_ID)
49    if type_name in _TYPE_NAME_TO_DECODER:
50      decoder = _TYPE_NAME_TO_DECODER[type_name]
51      return decoder(d)
52    else:
53      raise TypeError("Invalid type %s.", type_name)
54
55
56_DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S.%f"
57
58
59def _json_encode_datetime(o):
60  """Json encode a datetime object.
61
62  Args:
63    o: a datetime object.
64
65  Returns:
66    A dict of json primitives.
67  """
68  return {"isostr": o.strftime(_DATETIME_FORMAT)}
69
70
71def _json_decode_datetime(d):
72  """Converts a dict of json primitives to a datetime object."""
73  return datetime.datetime.strptime(d["isostr"], _DATETIME_FORMAT)
74
75
76def _register_json_primitive(object_type, encoder, decoder):
77  """Extend what MR can json serialize.
78
79  Args:
80    object_type: type of the object.
81    encoder: a function that takes in an object and returns a dict of
82       json primitives.
83    decoder: inverse function of encoder.
84  """
85  global _TYPE_TO_ENCODER
86  global _TYPE_NAME_TO_DECODER
87  if object_type not in _TYPE_TO_ENCODER:
88    _TYPE_TO_ENCODER[object_type] = encoder
89    _TYPE_NAME_TO_DECODER[object_type.__name__] = decoder
90
91
92_TYPE_TO_ENCODER = {}
93_TYPE_NAME_TO_DECODER = {}
94_register_json_primitive(datetime.datetime,
95                         _json_encode_datetime,
96                         _json_decode_datetime)
97
98# ndb.Key
99def _JsonEncodeKey(o):
100    """Json encode an ndb.Key object."""
101    return {'key_string': o.urlsafe()}
102
103def _JsonDecodeKey(d):
104    """Json decode a ndb.Key object."""
105    k_c = d['key_string']
106    if isinstance(k_c, (list, tuple)):
107        return ndb.Key(flat=k_c)
108    return ndb.Key(urlsafe=d['key_string'])
109
110_register_json_primitive(ndb.Key, _JsonEncodeKey, _JsonDecodeKey)
111
112
113class JsonMixin(object):
114  """Simple, stateless json utilities mixin.
115
116  Requires class to implement two methods:
117    to_json(self): convert data to json-compatible datastructure (dict,
118      list, strings, numbers)
119    @classmethod from_json(cls, json): load data from json-compatible structure.
120  """
121
122  def to_json_str(self):
123    """Convert data to json string representation.
124
125    Returns:
126      json representation as string.
127    """
128    _json = self.to_json()
129    try:
130      return json.dumps(_json, sort_keys=True, cls=JsonEncoder)
131    except:
132      logging.exception("Could not serialize JSON: %r", _json)
133      raise
134
135  @classmethod
136  def from_json_str(cls, json_str):
137    """Convert json string representation into class instance.
138
139    Args:
140      json_str: json representation as string.
141
142    Returns:
143      New instance of the class with data loaded from json string.
144    """
145    return cls.from_json(json.loads(json_str, cls=JsonDecoder))
146
147
148class JsonProperty(db.UnindexedProperty):
149  """Property type for storing json representation of data.
150
151  Requires data types to implement two methods:
152    to_json(self): convert data to json-compatible datastructure (dict,
153      list, strings, numbers)
154    @classmethod from_json(cls, json): load data from json-compatible structure.
155  """
156
157  def __init__(self, data_type, default=None, **kwargs):
158    """Constructor.
159
160    Args:
161      data_type: underlying data type as class.
162      default: default value for the property. The value is deep copied
163        fore each model instance.
164      **kwargs: remaining arguments.
165    """
166    kwargs["default"] = default
167    super(JsonProperty, self).__init__(**kwargs)
168    self.data_type = data_type
169
170  def get_value_for_datastore(self, model_instance):
171    """Gets value for datastore.
172
173    Args:
174      model_instance: instance of the model class.
175
176    Returns:
177      datastore-compatible value.
178    """
179    value = super(JsonProperty, self).get_value_for_datastore(model_instance)
180    if not value:
181      return None
182    json_value = value
183    if not isinstance(value, dict):
184      json_value = value.to_json()
185    if not json_value:
186      return None
187    return datastore_types.Text(json.dumps(
188        json_value, sort_keys=True, cls=JsonEncoder))
189
190  def make_value_from_datastore(self, value):
191    """Convert value from datastore representation.
192
193    Args:
194      value: datastore value.
195
196    Returns:
197      value to store in the model.
198    """
199
200    if value is None:
201      return None
202    _json = json.loads(value, cls=JsonDecoder)
203    if self.data_type == dict:
204      return _json
205    return self.data_type.from_json(_json)
206
207  def validate(self, value):
208    """Validate value.
209
210    Args:
211      value: model value.
212
213    Returns:
214      Whether the specified value is valid data type value.
215
216    Raises:
217      BadValueError: when value is not of self.data_type type.
218    """
219    if value is not None and not isinstance(value, self.data_type):
220      raise datastore_errors.BadValueError(
221          "Property %s must be convertible to a %s instance (%s)" %
222          (self.name, self.data_type, value))
223    return super(JsonProperty, self).validate(value)
224
225  def empty(self, value):
226    """Checks if value is empty.
227
228    Args:
229      value: model value.
230
231    Returns:
232      True passed value is empty.
233    """
234    return not value
235
236  def default_value(self):
237    """Create default model value.
238
239    If default option was specified, then it will be deeply copied.
240    None otherwise.
241
242    Returns:
243      default model value.
244    """
245    if self.default:
246      return copy.deepcopy(self.default)
247    else:
248      return None
249