hugr_cli/
hugr_io.rs

1//! Input/output arguments for the HUGR CLI.
2
3use clio::Input;
4use hugr::envelope::{EnvelopeConfig, EnvelopeError, read_envelope};
5use hugr::extension::ExtensionRegistry;
6use hugr::package::Package;
7use hugr::{Extension, Hugr};
8use std::io::{BufReader, Read};
9use std::path::PathBuf;
10
11use crate::CliError;
12
13/// Arguments for reading a HUGR input.
14#[derive(Debug, clap::Args)]
15pub struct HugrInputArgs {
16    /// Input file. Defaults to `-` for stdin.
17    #[arg(value_parser, default_value = "-", help_heading = "Input")]
18    pub input: Input,
19
20    /// No standard extensions.
21    #[arg(
22        long,
23        help_heading = "Input",
24        help = "Don't use standard extensions when validating hugrs. Prelude is still used."
25    )]
26    pub no_std: bool,
27    /// Extensions paths.
28    #[arg(
29        short,
30        long,
31        help_heading = "Input",
32        help = "Paths to serialised extensions to validate against."
33    )]
34    pub extensions: Vec<PathBuf>,
35    /// Read the input as a HUGR JSON file instead of an envelope.
36    ///
37    /// This is a legacy option for reading old HUGR files.
38    #[clap(long, help_heading = "Input")]
39    pub hugr_json: bool,
40}
41
42impl HugrInputArgs {
43    /// Read a hugr envelope from the input and return the package encoded
44    /// within.
45    ///
46    /// # Errors
47    ///
48    /// If [`HugrInputArgs::hugr_json`] is `true`, [`HugrInputArgs::get_hugr`] should be called instead as
49    /// reading the input as a package will fail.
50    pub fn get_package(&mut self) -> Result<Package, CliError> {
51        self.get_envelope().map(|(_, package)| package)
52    }
53
54    /// Read a hugr envelope from the input and return the envelope
55    /// configuration and the package encoded within.
56    ///
57    /// # Errors
58    ///
59    /// If [`HugrInputArgs::hugr_json`] is `true`, [`HugrInputArgs::get_hugr`] should be called instead as
60    /// reading the input as a package will fail.
61    pub fn get_envelope(&mut self) -> Result<(EnvelopeConfig, Package), CliError> {
62        let extensions = self.load_extensions()?;
63        let buffer = BufReader::new(&mut self.input);
64        read_envelope(buffer, &extensions).map_err(|e| match e {
65            EnvelopeError::MagicNumber { .. } => CliError::NotAnEnvelope,
66            _ => CliError::Envelope(e),
67        })
68    }
69    /// Read a hugr JSON file from the input.
70    ///
71    /// This is a legacy option for reading old HUGR JSON files when the
72    /// [`HugrInputArgs::hugr_json`] flag is used.
73    ///
74    /// For most cases, [`HugrInputArgs::get_package`] should be called instead.
75    pub fn get_hugr(&mut self) -> Result<Hugr, CliError> {
76        let extensions = self.load_extensions()?;
77        let mut buffer = BufReader::new(&mut self.input);
78
79        /// Wraps the hugr JSON so that it defines a valid envelope.
80        const PREPEND: &str = r#"HUGRiHJv?@{"modules": ["#;
81        const APPEND: &str = r#"],"extensions": []}"#;
82
83        let mut envelope = PREPEND.to_string();
84        buffer.read_to_string(&mut envelope)?;
85        envelope.push_str(APPEND);
86
87        let hugr = Hugr::load_str(envelope, Some(&extensions))?;
88        Ok(hugr)
89    }
90
91    /// Return a register with the selected extensions.
92    ///
93    /// This includes the standard extensions if [`HugrInputArgs::no_std`] is `false`,
94    /// and the extensions loaded from the paths in [`HugrInputArgs::extensions`].
95    pub fn load_extensions(&self) -> Result<ExtensionRegistry, CliError> {
96        let mut reg = if self.no_std {
97            hugr::extension::PRELUDE_REGISTRY.to_owned()
98        } else {
99            hugr::std_extensions::STD_REG.to_owned()
100        };
101
102        for ext in &self.extensions {
103            let f = std::fs::File::open(ext)?;
104            let ext: Extension = serde_json::from_reader(f)?;
105            reg.register_updated(ext);
106        }
107
108        Ok(reg)
109    }
110}