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

tinymist_task/compute/
query.rs
