nominatim/utils/json_writer.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2023 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Streaming JSON encoder.
   9 """
  10 from typing import Any, TypeVar, Optional, Callable
  11 import io
  12 try:
  13     import ujson as json
  14 except ModuleNotFoundError:
  15     import json # type: ignore[no-redef]
  16
  17 T = TypeVar('T') # pylint: disable=invalid-name
  18
  19 class JsonWriter:
  20     """ JSON encoder that renders the output directly into an output
  21         stream. This is a very simple writer which produces JSON in a
  22         compact as possible form.
  23
  24         The writer does not check for syntactic correctness. It is the
  25         responsibility of the caller to call the write functions in an
  26         order that produces correct JSON.
  27
  28         All functions return the writer object itself so that function
  29         calls can be chained.
  30     """
  31
  32     def __init__(self) -> None:
  33         self.data = io.StringIO()
  34         self.pending = ''
  35
  36
  37     def __call__(self) -> str:
  38         """ Return the rendered JSON content as a string.
  39             The writer remains usable after calling this function.
  40         """
  41         if self.pending:
  42             assert self.pending in (']', '}')
  43             self.data.write(self.pending)
  44             self.pending = ''
  45         return self.data.getvalue()
  46
  47
  48     def start_object(self) -> 'JsonWriter':
  49         """ Write the open bracket of a JSON object.
  50         """
  51         if self.pending:
  52             self.data.write(self.pending)
  53         self.pending = '{'
  54         return self
  55
  56
  57     def end_object(self) -> 'JsonWriter':
  58         """ Write the closing bracket of a JSON object.
  59         """
  60         assert self.pending in (',', '{', '')
  61         if self.pending == '{':
  62             self.data.write(self.pending)
  63         self.pending = '}'
  64         return self
  65
  66
  67     def start_array(self) -> 'JsonWriter':
  68         """ Write the opening bracket of a JSON array.
  69         """
  70         if self.pending:
  71             self.data.write(self.pending)
  72         self.pending = '['
  73         return self
  74
  75
  76     def end_array(self) -> 'JsonWriter':
  77         """ Write the closing bracket of a JSON array.
  78         """
  79         assert self.pending in (',', '[', '')
  80         if self.pending == '[':
  81             self.data.write(self.pending)
  82         self.pending = ']'
  83         return self
  84
  85
  86     def key(self, name: str) -> 'JsonWriter':
  87         """ Write the key string of a JSON object.
  88         """
  89         assert self.pending
  90         self.data.write(self.pending)
  91         self.data.write(json.dumps(name, ensure_ascii=False))
  92         self.pending = ':'
  93         return self
  94
  95
  96     def value(self, value: Any) -> 'JsonWriter':
  97         """ Write out a value as JSON. The function uses the json.dumps()
  98             function for encoding the JSON. Thus any value that can be
  99             encoded by that function is permissible here.
 100         """
 101         return self.raw(json.dumps(value, ensure_ascii=False))
 102
 103
 104     def next(self) -> 'JsonWriter':
 105         """ Write out a delimiter comma between JSON object or array elements.
 106         """
 107         if self.pending:
 108             self.data.write(self.pending)
 109         self.pending = ','
 110         return self
 111
 112
 113     def raw(self, raw_json: str) -> 'JsonWriter':
 114         """ Write out the given value as is. This function is useful if
 115             a value is already available in JSON format.
 116         """
 117         if self.pending:
 118             self.data.write(self.pending)
 119             self.pending = ''
 120         self.data.write(raw_json)
 121         return self
 122
 123
 124     def keyval(self, key: str, value: Any) -> 'JsonWriter':
 125         """ Write out an object element with the given key and value.
 126             This is a shortcut for calling 'key()', 'value()' and 'next()'.
 127         """
 128         self.key(key)
 129         self.value(value)
 130         return self.next()
 131
 132
 133     def keyval_not_none(self, key: str, value: Optional[T],
 134                         transform: Optional[Callable[[T], Any]] = None) -> 'JsonWriter':
 135         """ Write out an object element only if the value is not None.
 136             If 'transform' is given, it must be a function that takes the
 137             value type and returns a JSON encodable type. The transform
 138             function will be called before the value is written out.
 139         """
 140         if value is not None:
 141             self.key(key)
 142             self.value(transform(value) if transform else value)
 143             self.next()
 144         return self