partiql_common/
pretty.rs

1use pretty::{Arena, DocAllocator, DocBuilder, Pretty};
2use std::io;
3use std::io::Write;
4use std::string::FromUtf8Error;
5use thiserror::Error;
6
7pub const PRETTY_INDENT_MINOR_NEST: isize = 2;
8pub const PRETTY_INDENT_SUBORDINATE_CLAUSE_NEST: isize = 6;
9
10#[derive(Debug, Error)]
11#[non_exhaustive]
12pub enum ToPrettyError {
13    #[error("IO error: `{0}`")]
14    IoError(#[from] std::io::Error),
15
16    #[error("FromUtf8Error: `{0}`")]
17    FromUtf8Error(#[from] FromUtf8Error),
18}
19
20pub type ToPrettyResult<T> = Result<T, ToPrettyError>;
21
22pub trait ToPretty {
23    /// Pretty-prints to a `String`.
24    fn to_pretty_string(&self, width: usize) -> ToPrettyResult<String> {
25        let mut out = Vec::new();
26        self.to_pretty(width, &mut out)?;
27        Ok(String::from_utf8(out)?)
28    }
29
30    /// Pretty-prints to a `std::io::Write` object.
31    fn to_pretty<W>(&self, width: usize, out: &mut W) -> ToPrettyResult<()>
32    where
33        W: ?Sized + io::Write;
34}
35
36impl<T> ToPretty for T
37where
38    T: PrettyDoc,
39{
40    fn to_pretty<W>(&self, width: usize, out: &mut W) -> ToPrettyResult<()>
41    where
42        W: ?Sized + Write,
43    {
44        let arena = Arena::new();
45        let DocBuilder(_, doc) = self.pretty_doc::<_, ()>(&arena);
46        Ok(doc.render(width, out)?)
47    }
48}
49
50pub trait PrettyDoc {
51    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
52    where
53        D: DocAllocator<'b, A>,
54        D::Doc: Clone,
55        A: Clone;
56}
57
58impl<T> PrettyDoc for &T
59where
60    T: PrettyDoc,
61{
62    #[inline]
63    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
64    where
65        D: DocAllocator<'b, A>,
66        D::Doc: Clone,
67        A: Clone,
68    {
69        (*self).pretty_doc(arena)
70    }
71}
72
73impl<T> PrettyDoc for Box<T>
74where
75    T: PrettyDoc,
76{
77    #[inline]
78    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
79    where
80        D: DocAllocator<'b, A>,
81        D::Doc: Clone,
82        A: Clone,
83    {
84        self.as_ref().pretty_doc(arena)
85    }
86}
87
88impl PrettyDoc for str {
89    #[inline]
90    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
91    where
92        D: DocAllocator<'b, A>,
93        D::Doc: Clone,
94        A: Clone,
95    {
96        arena.concat(["'", self, "'"])
97    }
98}
99
100impl PrettyDoc for String {
101    #[inline]
102    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
103    where
104        D: DocAllocator<'b, A>,
105        D::Doc: Clone,
106        A: Clone,
107    {
108        arena.concat(["'", self, "'"])
109    }
110}
111
112impl PrettyDoc for Vec<u8> {
113    #[inline]
114    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
115    where
116        D: DocAllocator<'b, A>,
117        D::Doc: Clone,
118        A: Clone,
119    {
120        let y = String::from_utf8_lossy(self.as_slice());
121        arena.text(y)
122    }
123}
124
125impl PrettyDoc for rust_decimal::Decimal {
126    #[inline]
127    fn pretty_doc<'b, D, A>(&'b self, arena: &'b D) -> DocBuilder<'b, D, A>
128    where
129        D: DocAllocator<'b, A>,
130        D::Doc: Clone,
131        A: Clone,
132    {
133        arena.text(self.to_string())
134    }
135}
136
137#[inline]
138pub fn pretty_prefixed_doc<'b, E, D, A>(
139    annot: &'static str,
140    doc: E,
141    arena: &'b D,
142) -> DocBuilder<'b, D, A>
143where
144    E: Pretty<'b, D, A>,
145    D: DocAllocator<'b, A>,
146    D::Doc: Clone,
147    A: Clone,
148{
149    arena.text(annot).append(arena.space()).append(doc).group()
150}
151
152#[inline]
153pub fn pretty_surrounded<'b, P, D, A>(
154    inner: &'b P,
155    start: &'static str,
156    end: &'static str,
157    arena: &'b D,
158) -> DocBuilder<'b, D, A>
159where
160    P: PrettyDoc + 'b,
161    D: DocAllocator<'b, A>,
162    D::Doc: Clone,
163    A: Clone,
164{
165    pretty_surrounded_doc(inner.pretty_doc(arena), start, end, arena)
166}
167
168#[inline]
169pub fn pretty_surrounded_doc<'b, E, D, A>(
170    doc: E,
171    start: &'static str,
172    end: &'static str,
173    arena: &'b D,
174) -> DocBuilder<'b, D, A>
175where
176    E: Pretty<'b, D, A>,
177    D: DocAllocator<'b, A>,
178    D::Doc: Clone,
179    A: Clone,
180{
181    arena
182        .text(start)
183        .append(doc)
184        .append(arena.text(end))
185        .group()
186}
187
188#[inline]
189pub fn pretty_parenthesized_doc<'b, E, D, A>(doc: E, arena: &'b D) -> DocBuilder<'b, D, A>
190where
191    E: Pretty<'b, D, A>,
192    D: DocAllocator<'b, A>,
193    D::Doc: Clone,
194    A: Clone,
195{
196    pretty_surrounded_doc(doc, "(", ")", arena)
197}
198
199#[inline]
200pub fn pretty_seq_doc<'i, 'b, I, E, D, A>(
201    seq: I,
202    start: &'static str,
203    qualifier: Option<E>,
204    end: &'static str,
205    sep: &'static str,
206    nest: isize,
207    arena: &'b D,
208) -> DocBuilder<'b, D, A>
209where
210    E: Pretty<'b, D, A>,
211    I: IntoIterator<Item = E>,
212    D: DocAllocator<'b, A>,
213    D::Doc: Clone,
214    A: Clone,
215{
216    let start = arena.text(start);
217    let end = arena.text(end);
218    let sep = arena.text(sep).append(arena.line());
219    let start = if let Some(qual) = qualifier {
220        start.append(arena.space()).append(qual)
221    } else {
222        start
223    };
224    let body = arena
225        .line()
226        .append(arena.intersperse(seq, sep))
227        .append(arena.line())
228        .group();
229    start.append(body.nest(nest)).append(end).group()
230}
231
232#[inline]
233pub fn pretty_seq<'i, 'b, I, P, D, A>(
234    list: I,
235    start: &'static str,
236    end: &'static str,
237    sep: &'static str,
238    nest: isize,
239    arena: &'b D,
240) -> DocBuilder<'b, D, A>
241where
242    I: IntoIterator<Item = &'b P>,
243    P: PrettyDoc + 'b,
244    D: DocAllocator<'b, A>,
245    D::Doc: Clone,
246    A: Clone,
247{
248    let seq = list.into_iter().map(|l| l.pretty_doc(arena));
249    pretty_seq_doc(seq, start, None, end, sep, nest, arena)
250}
251
252#[inline]
253pub fn pretty_list<'b, I, P, D, A>(list: I, nest: isize, arena: &'b D) -> DocBuilder<'b, D, A>
254where
255    I: IntoIterator<Item = &'b P>,
256    P: PrettyDoc + 'b,
257    D: DocAllocator<'b, A>,
258    D::Doc: Clone,
259    A: Clone,
260{
261    let sep = arena.text(",").append(arena.softline());
262    pretty_seperated(sep, list, nest, arena)
263}
264
265#[inline]
266pub fn pretty_doc_list<'b, I, D, A>(list: I, nest: isize, arena: &'b D) -> DocBuilder<'b, D, A>
267where
268    I: IntoIterator<Item = DocBuilder<'b, D, A>>,
269    D: DocAllocator<'b, A>,
270    D::Doc: Clone,
271    A: Clone,
272{
273    let sep = arena.text(",").append(arena.softline());
274    pretty_seperated_doc(sep, list, nest, arena)
275}
276
277#[inline]
278pub fn pretty_seperated<'b, I, E, P, D, A>(
279    sep: E,
280    list: I,
281    nest: isize,
282    arena: &'b D,
283) -> DocBuilder<'b, D, A>
284where
285    I: IntoIterator<Item = &'b P>,
286    E: Pretty<'b, D, A>,
287    P: PrettyDoc + 'b,
288    D: DocAllocator<'b, A>,
289    D::Doc: Clone,
290    A: Clone,
291{
292    let list = list.into_iter().map(|l| l.pretty_doc(arena));
293    pretty_seperated_doc(sep, list, nest, arena)
294}
295
296#[inline]
297pub fn pretty_seperated_doc<'b, I, E, D, A>(
298    sep: E,
299    list: I,
300    nest: isize,
301    arena: &'b D,
302) -> DocBuilder<'b, D, A>
303where
304    I: IntoIterator<Item = DocBuilder<'b, D, A>>,
305    E: Pretty<'b, D, A>,
306    D: DocAllocator<'b, A>,
307    D::Doc: Clone,
308    A: Clone,
309{
310    let sep = sep.pretty(arena);
311    arena.intersperse(list, sep).nest(nest).group()
312}