feat: support vscode tasks for exporting query and pdfpc (#490)

* feat: support vscode tasks for exporting query and pdfpc

* test: update snapshot

crates/tinymist-query/src/lib.rs

@@ -153,6 +153,15 @@ mod polymorphic {
         Html {},
         Markdown {},
         Text {},
+        Query {
+            format: String,
+            output_extension: Option<String>,
+            strict: bool,
+            selector: String,
+            field: Option<String>,
+            one: bool,
+            pretty: bool,
+        },
         Svg {
             page: PageSelection,
         },
@@ -180,6 +189,11 @@ mod polymorphic {
                 Self::Text { .. } => "txt",
                 Self::Svg { .. } => "svg",
                 Self::Png { .. } => "png",
+                Self::Query {
+                    format,
+                    output_extension,
+                    ..
+                } => output_extension.as_deref().unwrap_or(format),
             }
         }
     }

crates/tinymist/Cargo.toml

@@ -31,6 +31,7 @@ env_logger.workspace = true
 log.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+serde_yaml.workspace = true
 parking_lot.workspace = true
 paste.workspace = true
 

crates/tinymist/src/cmd.rs

@@ -26,6 +26,18 @@ struct ExportOpts {
     page: PageSelection,
 }
 
+#[derive(Debug, Clone, Default, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct QueryOpts {
+    format: String,
+    output_extension: Option<String>,
+    strict: Option<bool>,
+    pretty: Option<bool>,
+    selector: String,
+    field: Option<String>,
+    one: Option<bool>,
+}
+
 /// Here are implemented the handlers for each command.
 impl LanguageState {
     /// Export the current document as PDF file(s).
@@ -59,6 +71,24 @@ impl LanguageState {
         self.export(req_id, ExportKind::Text {}, args)
     }
 
+    /// Query the current document and export the result as JSON file(s).
+    pub fn export_query(&mut self, req_id: RequestId, mut args: Vec<JsonValue>) -> ScheduledResult {
+        let opts = get_arg_or_default!(args[1] as QueryOpts);
+        self.export(
+            req_id,
+            ExportKind::Query {
+                format: opts.format,
+                output_extension: opts.output_extension,
+                strict: opts.strict.unwrap_or(true),
+                selector: opts.selector,
+                field: opts.field,
+                pretty: opts.pretty.unwrap_or(true),
+                one: opts.one.unwrap_or(false),
+            },
+            args,
+        )
+    }
+
     /// Export the current document as Svg file(s).
     pub fn export_svg(&mut self, req_id: RequestId, mut args: Vec<JsonValue>) -> ScheduledResult {
         let opts = get_arg_or_default!(args[1] as ExportOpts);

crates/tinymist/src/server.rs

@@ -257,6 +257,7 @@ impl LanguageState {
             .with_command_("tinymist.exportText", State::export_text)
             .with_command_("tinymist.exportHtml", State::export_html)
             .with_command_("tinymist.exportMarkdown", State::export_markdown)
+            .with_command_("tinymist.exportQuery", State::export_query)
             .with_command("tinymist.doClearCache", State::clear_cache)
             .with_command("tinymist.pinMain", State::pin_document)
             .with_command("tinymist.focusMain", State::focus_document)

crates/tinymist/src/task/export.rs

@@ -1,5 +1,6 @@
 //! The actor that handles various document export, like PDF and SVG export.
 
+use std::ops::Deref;
 use std::str::FromStr;
 use std::{path::PathBuf, sync::Arc};
 
@@ -8,6 +9,7 @@ use once_cell::sync::Lazy;
 use tinymist_query::{ExportKind, PageSelection};
 use tokio::sync::mpsc;
 use typlite::Typlite;
+use typst::foundations::IntoValue;
 use typst::{
     foundations::Smart,
     layout::{Abs, Frame},
@@ -213,6 +215,39 @@ impl ExportConfig {
                     // todo: timestamp world.now()
                     typst_pdf::pdf(doc, Smart::Auto, timestamp)
                 }
+                Query {
+                    format,
+                    output_extension: _,
+                    strict,
+                    selector,
+                    field,
+                    one,
+                    pretty,
+                } => {
+                    let elements =
+                        typst_ts_compiler::query::retrieve(artifact.world.deref(), &selector, doc)
+                            .map_err(|e| anyhow::anyhow!("failed to retrieve: {e}"))?;
+                    if one && elements.len() != 1 {
+                        bail!("expected exactly one element, found {}", elements.len());
+                    }
+
+                    let mapped: Vec<_> = elements
+                        .into_iter()
+                        .filter_map(|c| match &field {
+                            Some(field) => c.get_by_name(field),
+                            _ => Some(c.into_value()),
+                        })
+                        .collect();
+
+                    if one {
+                        let Some(value) = mapped.first() else {
+                            bail!("no such field found for element");
+                        };
+                        serialize(value, &format, strict, pretty).map(String::into_bytes)?
+                    } else {
+                        serialize(&mapped, &format, strict, pretty).map(String::into_bytes)?
+                    }
+                }
                 Html {} => typst_ts_svg_exporter::render_svg_html(doc).into_bytes(),
                 Text {} => format!("{}", FullTextDigest(doc.clone())).into_bytes(),
                 Markdown {} => {
@@ -355,6 +390,39 @@ fn convert_datetime(date_time: chrono::DateTime<chrono::Utc>) -> Option<TypstDat
     )
 }
 
+/// Serialize data to the output format.
+fn serialize(
+    data: &impl serde::Serialize,
+    format: &str,
+    strict: bool,
+    pretty: bool,
+) -> anyhow::Result<String> {
+    Ok(match format {
+        "json" if pretty => serde_json::to_string_pretty(data)?,
+        "json" => serde_json::to_string(data)?,
+        "yaml" => serde_yaml::to_string(&data)?,
+        format if format == "txt" || !strict => {
+            use serde_json::Value::*;
+            let value = serde_json::to_value(data)?;
+            match value {
+                String(s) => s,
+                _ => {
+                    let kind = match value {
+                        Null => "null",
+                        Bool(_) => "boolean",
+                        Number(_) => "number",
+                        String(_) => "string",
+                        Array(_) => "array",
+                        Object(_) => "object",
+                    };
+                    bail!("expected a string value for format: {format}, got {kind}")
+                }
+            }
+        }
+        _ => bail!("unsupported format for query: {format}"),
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;