tinymist_task/compute/
query.rs

1//! The computation for document query.
2
3use std::sync::Arc;
4
5use comemo::Track;
6use ecow::EcoString;
7use tinymist_std::error::prelude::*;
8use tinymist_std::typst::TypstDocument;
9use tinymist_world::{CompilerFeat, ExportComputation, WorldComputeGraph};
10use typst::World;
11use typst::diag::{SourceResult, StrResult};
12use typst::engine::Sink;
13use typst::foundations::{Content, Context, IntoValue, LocatableSelector, Output, Scope, Value};
14use typst::model::Document;
15use typst::routines::SpanMode;
16use typst::syntax::Span;
17use typst::syntax::SyntaxMode;
18use typst_eval::eval_string;
19
20use crate::QueryTask;
21
22/// The computation for document query.
23pub struct DocumentQuery;
24
25impl DocumentQuery {
26    // todo: query exporter
27    /// Retrieve the matches for the selector.
28    pub fn retrieve<D: Document + Output>(
29        world: &dyn World,
30        selector: &str,
31        document: &D,
32    ) -> StrResult<Vec<Content>> {
33        let selector = eval_string(
34            world.track(),
35            world.library(),
36            Sink::new().track_mut(),
37            document.introspector().track(),
38            Context::none().track(),
39            selector,
40            SpanMode::Uniform(Span::detached()),
41            SyntaxMode::Code,
42            Scope::default(),
43        )
44        .map_err(|errors| {
45            let mut message = EcoString::from("failed to evaluate selector");
46            for (i, error) in errors.into_iter().enumerate() {
47                message.push_str(if i == 0 { ": " } else { ", " });
48                message.push_str(&error.message);
49            }
50            message
51        })?
52        .cast::<LocatableSelector>()
53        .map_err(|e| EcoString::from(format!("failed to cast: {}", e.message())))?;
54
55        Ok(document
56            .introspector()
57            .query(&selector.0)
58            .into_iter()
59            .collect::<Vec<_>>())
60    }
61
62    fn run_inner<F: CompilerFeat, D: Document + Output>(
63        g: &Arc<WorldComputeGraph<F>>,
64        doc: &Arc<D>,
65        config: &QueryTask,
66    ) -> Result<Vec<Value>> {
67        let selector = &config.selector;
68        let elements = Self::retrieve(&g.snap.world, selector, doc.as_ref())
69            .map_err(|e| anyhow::anyhow!("failed to retrieve: {e}"))?;
70        if config.one && elements.len() != 1 {
71            bail!("expected exactly one element, found {}", elements.len());
72        }
73
74        Ok(elements
75            .into_iter()
76            .filter_map(|c| match &config.field {
77                Some(field) => c.get_by_name(field).ok(),
78                _ => Some(c.into_value()),
79            })
80            .collect())
81    }
82
83    /// Queries the document and returns the result as a value.
84    pub fn doc_get_as_value<F: CompilerFeat>(
85        g: &Arc<WorldComputeGraph<F>>,
86        doc: &TypstDocument,
87        config: &QueryTask,
88    ) -> Result<serde_json::Value> {
89        match doc {
90            TypstDocument::Paged(doc) => Self::get_as_value(g, doc, config),
91            TypstDocument::Html(doc) => Self::get_as_value(g, doc, config),
92        }
93    }
94
95    /// Queries the document and returns the result as a value.
96    pub fn get_as_value<F: CompilerFeat, D: Document + Output>(
97        g: &Arc<WorldComputeGraph<F>>,
98        doc: &Arc<D>,
99        config: &QueryTask,
100    ) -> Result<serde_json::Value> {
101        let mapped = Self::run_inner(g, doc, config)?;
102
103        let res = if config.one {
104            let Some(value) = mapped.first() else {
105                bail!("no such field found for element");
106            };
107            serde_json::to_value(value)
108        } else {
109            serde_json::to_value(&mapped)
110        };
111
112        res.context("failed to serialize")
113    }
114}
115
116impl<F: CompilerFeat, D: Document + Output> ExportComputation<F, D> for DocumentQuery {
117    type Output = SourceResult<String>;
118    type Config = QueryTask;
119
120    fn run(
121        g: &Arc<WorldComputeGraph<F>>,
122        doc: &Arc<D>,
123        config: &QueryTask,
124    ) -> Result<SourceResult<String>> {
125        let pretty = false;
126        let mapped = Self::run_inner(g, doc, config)?;
127
128        let res = if config.one {
129            let Some(value) = mapped.first() else {
130                bail!("no such field found for element");
131            };
132            serialize(value, &config.format, pretty)
133        } else {
134            serialize(&mapped, &config.format, pretty)
135        };
136
137        res.map(Ok)
138    }
139}
140
141/// Serialize data to the output format.
142fn serialize(data: &impl serde::Serialize, format: &str, pretty: bool) -> Result<String> {
143    Ok(match format {
144        "json" if pretty => serde_json::to_string_pretty(data).context("serialize query")?,
145        "json" => serde_json::to_string(data).context("serialize query")?,
146        "yaml" => serde_yaml::to_string(&data).context_ut("serialize query")?,
147        "txt" => {
148            use serde_json::Value::*;
149            let value = serde_json::to_value(data).context("serialize query")?;
150            match value {
151                String(s) => s,
152                _ => {
153                    let kind = match value {
154                        Null => "null",
155                        Bool(_) => "boolean",
156                        Number(_) => "number",
157                        String(_) => "string",
158                        Array(_) => "array",
159                        Object(_) => "object",
160                    };
161                    bail!("expected a string value for format: {format}, got {kind}")
162                }
163            }
164        }
165        _ => bail!("unsupported format for query: {format}"),
166    })
167}
tinymist_task/compute/query.rs

tinymist_task/compute/
query.rs
