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

tinymist_task/compute/
query.rs
