]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/utils/json_writer.py
make sure PHP and Python reverse code does the same
[nominatim.git] / 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 float(self, value: float, precision: int) -> 'JsonWriter':
105         """ Write out a float value with the given precision.
106         """
107         return self.raw(f"{value:0.{precision}f}")
108
109     def next(self) -> 'JsonWriter':
110         """ Write out a delimiter comma between JSON object or array elements.
111         """
112         if self.pending:
113             self.data.write(self.pending)
114         self.pending = ','
115         return self
116
117
118     def raw(self, raw_json: str) -> 'JsonWriter':
119         """ Write out the given value as is. This function is useful if
120             a value is already available in JSON format.
121         """
122         if self.pending:
123             self.data.write(self.pending)
124             self.pending = ''
125         self.data.write(raw_json)
126         return self
127
128
129     def keyval(self, key: str, value: Any) -> 'JsonWriter':
130         """ Write out an object element with the given key and value.
131             This is a shortcut for calling 'key()', 'value()' and 'next()'.
132         """
133         self.key(key)
134         self.value(value)
135         return self.next()
136
137
138     def keyval_not_none(self, key: str, value: Optional[T],
139                         transform: Optional[Callable[[T], Any]] = None) -> 'JsonWriter':
140         """ Write out an object element only if the value is not None.
141             If 'transform' is given, it must be a function that takes the
142             value type and returns a JSON encodable type. The transform
143             function will be called before the value is written out.
144         """
145         if value is not None:
146             self.key(key)
147             self.value(transform(value) if transform else value)
148             self.next()
149         return self