src/nominatim_api/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) 2024 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')
  18
  19
  20 class JsonWriter:
  21     """ JSON encoder that renders the output directly into an output
  22         stream. This is a very simple writer which produces JSON in a
  23         compact as possible form.
  24
  25         The writer does not check for syntactic correctness. It is the
  26         responsibility of the caller to call the write functions in an
  27         order that produces correct JSON.
  28
  29         All functions return the writer object itself so that function
  30         calls can be chained.
  31     """
  32
  33     def __init__(self) -> None:
  34         self.data = io.StringIO()
  35         self.pending = ''
  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     def start_object(self) -> 'JsonWriter':
  48         """ Write the open bracket of a JSON object.
  49         """
  50         if self.pending:
  51             self.data.write(self.pending)
  52         self.pending = '{'
  53         return self
  54
  55     def end_object(self) -> 'JsonWriter':
  56         """ Write the closing bracket of a JSON object.
  57         """
  58         assert self.pending in (',', '{', '')
  59         if self.pending == '{':
  60             self.data.write(self.pending)
  61         self.pending = '}'
  62         return self
  63
  64     def start_array(self) -> 'JsonWriter':
  65         """ Write the opening bracket of a JSON array.
  66         """
  67         if self.pending:
  68             self.data.write(self.pending)
  69         self.pending = '['
  70         return self
  71
  72     def end_array(self) -> 'JsonWriter':
  73         """ Write the closing bracket of a JSON array.
  74         """
  75         assert self.pending in (',', '[', ']', ')', '')
  76         if self.pending not in (',', ''):
  77             self.data.write(self.pending)
  78         self.pending = ']'
  79         return self
  80
  81     def key(self, name: str) -> 'JsonWriter':
  82         """ Write the key string of a JSON object.
  83         """
  84         assert self.pending
  85         self.data.write(self.pending)
  86         self.data.write(json.dumps(name, ensure_ascii=False))
  87         self.pending = ':'
  88         return self
  89
  90     def value(self, value: Any) -> 'JsonWriter':
  91         """ Write out a value as JSON. The function uses the json.dumps()
  92             function for encoding the JSON. Thus any value that can be
  93             encoded by that function is permissible here.
  94         """
  95         return self.raw(json.dumps(value, ensure_ascii=False))
  96
  97     def float(self, value: float, precision: int) -> 'JsonWriter':
  98         """ Write out a float value with the given precision.
  99         """
 100         return self.raw(f"{value:0.{precision}f}")
 101
 102     def next(self) -> 'JsonWriter':
 103         """ Write out a delimiter comma between JSON object or array elements.
 104         """
 105         if self.pending:
 106             self.data.write(self.pending)
 107         self.pending = ','
 108         return self
 109
 110     def raw(self, raw_json: str) -> 'JsonWriter':
 111         """ Write out the given value as is. This function is useful if
 112             a value is already available in JSON format.
 113         """
 114         if self.pending:
 115             self.data.write(self.pending)
 116             self.pending = ''
 117         self.data.write(raw_json)
 118         return self
 119
 120     def keyval(self, key: str, value: Any) -> 'JsonWriter':
 121         """ Write out an object element with the given key and value.
 122             This is a shortcut for calling 'key()', 'value()' and 'next()'.
 123         """
 124         self.key(key)
 125         self.value(value)
 126         return self.next()
 127
 128     def keyval_not_none(self, key: str, value: Optional[T],
 129                         transform: Optional[Callable[[T], Any]] = None) -> 'JsonWriter':
 130         """ Write out an object element only if the value is not None.
 131             If 'transform' is given, it must be a function that takes the
 132             value type and returns a JSON encodable type. The transform
 133             function will be called before the value is written out.
 134         """
 135         if value is not None:
 136             self.key(key)
 137             self.value(transform(value) if transform else value)
 138             self.next()
 139         return self