fauna.encoding.encoder

  1import base64
  2from datetime import datetime, date
  3from typing import Any, Optional, List, Union
  4
  5from fauna.query.models import DocumentReference, Module, Document, NamedDocument, NamedDocumentReference, NullDocument, \
  6  StreamToken
  7from fauna.query.query_builder import Query, Fragment, LiteralFragment, ValueFragment
  8
  9_RESERVED_TAGS = [
 10    "@date",
 11    "@doc",
 12    "@double",
 13    "@int",
 14    "@long",
 15    "@mod",
 16    "@object",
 17    "@ref",
 18    "@set",
 19    "@time",
 20]
 21
 22
 23class FaunaEncoder:
 24  """Supports the following types:
 25
 26    +-------------------------------+---------------+
 27    | Python                        | Fauna Tags    |
 28    +===============================+===============+
 29    | dict                          | @object       |
 30    +-------------------------------+---------------+
 31    | list, tuple                   | array         |
 32    +-------------------------------+---------------+
 33    | str                           | string        |
 34    +-------------------------------+---------------+
 35    | int 32-bit signed             | @int          |
 36    +-------------------------------+---------------+
 37    | int 64-bit signed             | @long         |
 38    +-------------------------------+---------------+
 39    | float                         | @double       |
 40    +-------------------------------+---------------+
 41    | datetime.datetime             | @time         |
 42    +-------------------------------+---------------+
 43    | datetime.date                 | @date         |
 44    +-------------------------------+---------------+
 45    | True                          | True          |
 46    +-------------------------------+---------------+
 47    | False                         | False         |
 48    +-------------------------------+---------------+
 49    | None                          | None          |
 50    +-------------------------------+---------------+
 51    | bytes / bytearray             | @bytes        |
 52    +-------------------------------+---------------+
 53    | *Document                     | @ref          |
 54    +-------------------------------+---------------+
 55    | *DocumentReference            | @ref          |
 56    +-------------------------------+---------------+
 57    | Module                        | @mod          |
 58    +-------------------------------+---------------+
 59    | Query                         | fql           |
 60    +-------------------------------+---------------+
 61    | ValueFragment                 | value         |
 62    +-------------------------------+---------------+
 63    | TemplateFragment              | string        |
 64    +-------------------------------+---------------+
 65    | StreamToken                   | string        |
 66    +-------------------------------+---------------+
 67
 68    """
 69
 70  @staticmethod
 71  def encode(obj: Any) -> Any:
 72    """Encodes supported objects into the tagged format.
 73
 74        Examples:
 75            - Up to 32-bit ints encode to { "@int": "..." }
 76            - Up to 64-bit ints encode to { "@long": "..." }
 77            - Floats encode to { "@double": "..." }
 78            - datetime encodes to { "@time": "..." }
 79            - date encodes to { "@date": "..." }
 80            - DocumentReference encodes to { "@doc": "..." }
 81            - Module encodes to { "@mod": "..." }
 82            - Query encodes to { "fql": [...] }
 83            - ValueFragment encodes to { "value": <encoded_val> }
 84            - LiteralFragment encodes to a string
 85            - StreamToken encodes to a string
 86
 87        :raises ValueError: If value cannot be encoded, cannot be encoded safely, or there's a circular reference.
 88        :param obj: the object to decode
 89        """
 90    return FaunaEncoder._encode(obj)
 91
 92  @staticmethod
 93  def from_int(obj: int):
 94    if -2**31 <= obj <= 2**31 - 1:
 95      return {"@int": repr(obj)}
 96    elif -2**63 <= obj <= 2**63 - 1:
 97      return {"@long": repr(obj)}
 98    else:
 99      raise ValueError("Precision loss when converting int to Fauna type")
100
101  @staticmethod
102  def from_bool(obj: bool):
103    return obj
104
105  @staticmethod
106  def from_float(obj: float):
107    return {"@double": repr(obj)}
108
109  @staticmethod
110  def from_str(obj: str):
111    return obj
112
113  @staticmethod
114  def from_datetime(obj: datetime):
115    if obj.utcoffset() is None:
116      raise ValueError("datetimes must be timezone-aware")
117
118    return {"@time": obj.isoformat(sep="T")}
119
120  @staticmethod
121  def from_date(obj: date):
122    return {"@date": obj.isoformat()}
123
124  @staticmethod
125  def from_bytes(obj: Union[bytearray, bytes]):
126    return {"@bytes": base64.b64encode(obj).decode('ascii')}
127
128  @staticmethod
129  def from_doc_ref(obj: DocumentReference):
130    return {"@ref": {"id": obj.id, "coll": FaunaEncoder.from_mod(obj.coll)}}
131
132  @staticmethod
133  def from_named_doc_ref(obj: NamedDocumentReference):
134    return {"@ref": {"name": obj.name, "coll": FaunaEncoder.from_mod(obj.coll)}}
135
136  @staticmethod
137  def from_mod(obj: Module):
138    return {"@mod": obj.name}
139
140  @staticmethod
141  def from_dict(obj: Any):
142    return {"@object": obj}
143
144  @staticmethod
145  def from_none():
146    return None
147
148  @staticmethod
149  def from_fragment(obj: Fragment):
150    if isinstance(obj, LiteralFragment):
151      return obj.get()
152    elif isinstance(obj, ValueFragment):
153      v = obj.get()
154      if isinstance(v, Query):
155        return FaunaEncoder.from_query_interpolation_builder(v)
156      else:
157        return {"value": FaunaEncoder.encode(v)}
158    else:
159      raise ValueError(f"Unknown fragment type: {type(obj)}")
160
161  @staticmethod
162  def from_query_interpolation_builder(obj: Query):
163    return {"fql": [FaunaEncoder.from_fragment(f) for f in obj.fragments]}
164
165  @staticmethod
166  def from_streamtoken(obj: StreamToken):
167    return {"@stream": obj.token}
168
169  @staticmethod
170  def _encode(o: Any, _markers: Optional[List] = None):
171    if _markers is None:
172      _markers = []
173
174    if isinstance(o, str):
175      return FaunaEncoder.from_str(o)
176    elif o is None:
177      return FaunaEncoder.from_none()
178    elif o is True:
179      return FaunaEncoder.from_bool(o)
180    elif o is False:
181      return FaunaEncoder.from_bool(o)
182    elif isinstance(o, int):
183      return FaunaEncoder.from_int(o)
184    elif isinstance(o, float):
185      return FaunaEncoder.from_float(o)
186    elif isinstance(o, Module):
187      return FaunaEncoder.from_mod(o)
188    elif isinstance(o, DocumentReference):
189      return FaunaEncoder.from_doc_ref(o)
190    elif isinstance(o, NamedDocumentReference):
191      return FaunaEncoder.from_named_doc_ref(o)
192    elif isinstance(o, datetime):
193      return FaunaEncoder.from_datetime(o)
194    elif isinstance(o, date):
195      return FaunaEncoder.from_date(o)
196    elif isinstance(o, bytearray) or isinstance(o, bytes):
197      return FaunaEncoder.from_bytes(o)
198    elif isinstance(o, Document):
199      return FaunaEncoder.from_doc_ref(DocumentReference(o.coll, o.id))
200    elif isinstance(o, NamedDocument):
201      return FaunaEncoder.from_named_doc_ref(
202          NamedDocumentReference(o.coll, o.name))
203    elif isinstance(o, NullDocument):
204      return FaunaEncoder.encode(o.ref)
205    elif isinstance(o, (list, tuple)):
206      return FaunaEncoder._encode_list(o, _markers)
207    elif isinstance(o, dict):
208      return FaunaEncoder._encode_dict(o, _markers)
209    elif isinstance(o, Query):
210      return FaunaEncoder.from_query_interpolation_builder(o)
211    elif isinstance(o, StreamToken):
212      return FaunaEncoder.from_streamtoken(o)
213    else:
214      raise ValueError(f"Object {o} of type {type(o)} cannot be encoded")
215
216  @staticmethod
217  def _encode_list(lst, markers):
218    _id = id(lst)
219    if _id in markers:
220      raise ValueError("Circular reference detected")
221
222    markers.append(id(lst))
223    res = [FaunaEncoder._encode(elem, markers) for elem in lst]
224    markers.pop()
225    return res
226
227  @staticmethod
228  def _encode_dict(dct, markers):
229    _id = id(dct)
230    if _id in markers:
231      raise ValueError("Circular reference detected")
232
233    markers.append(id(dct))
234    if any(i in _RESERVED_TAGS for i in dct.keys()):
235      res = {
236          "@object": {
237              k: FaunaEncoder._encode(v, markers) for k, v in dct.items()
238          }
239      }
240      markers.pop()
241      return res
242    else:
243      res = {k: FaunaEncoder._encode(v, markers) for k, v in dct.items()}
244      markers.pop()
245      return res