yaml_edit/nodes/

entry.rs

//! Entry API for [`Mapping`](crate::Mapping).
//!
//! Modelled on [`std::collections::BTreeMap::entry`]. Obtain an [`Entry`]
//! via [`Mapping::entry`](crate::Mapping::entry).
//!
//! ```rust
//! use std::str::FromStr;
//! use yaml_edit::Document;
//!
//! let doc = Document::from_str("name: Alice\n").unwrap();
//! let mapping = doc.as_mapping().unwrap();
//!
//! mapping.entry("age").or_insert(30);
//! mapping.entry("name").or_insert("Bob");
//!
//! assert_eq!(doc.to_string(), "name: Alice\nage: 30\n");
//! ```

use crate::as_yaml::{AsYaml, YamlNode};
use crate::nodes::mapping::{Mapping, MappingEntry};
use crate::nodes::sequence::Sequence;

/// A view into a single entry in a [`Mapping`], which may either be vacant
/// or occupied.
///
/// Obtained from [`Mapping::entry`](crate::Mapping::entry).
#[derive(Debug)]
pub enum Entry<'a, K: AsYaml> {
    /// The key is already present in the mapping.
    Occupied(OccupiedEntry<'a>),
    /// The key is not present in the mapping.
    Vacant(VacantEntry<'a, K>),
}

/// A view into an occupied entry in a [`Mapping`].
#[derive(Debug)]
pub struct OccupiedEntry<'a> {
    pub(crate) mapping: &'a Mapping,
    pub(crate) entry: MappingEntry,
}

/// A view into a vacant entry in a [`Mapping`].
#[derive(Debug)]
pub struct VacantEntry<'a, K: AsYaml> {
    pub(crate) mapping: &'a Mapping,
    pub(crate) key: K,
}

impl<'a, K: AsYaml> Entry<'a, K> {
    /// Ensure a value is in the entry by inserting `default` if vacant, and
    /// return the resulting value as a [`YamlNode`].
    pub fn or_insert(self, default: impl AsYaml) -> YamlNode {
        match self {
            Entry::Occupied(o) => o.into_value(),
            Entry::Vacant(v) => v.insert(default),
        }
    }

    /// Ensure a value is in the entry by inserting the result of `default` if
    /// vacant, and return the resulting value as a [`YamlNode`].
    ///
    /// `default` is only called if the entry is vacant.
    pub fn or_insert_with<V, F>(self, default: F) -> YamlNode
    where
        V: AsYaml,
        F: FnOnce() -> V,
    {
        match self {
            Entry::Occupied(o) => o.into_value(),
            Entry::Vacant(v) => v.insert(default()),
        }
    }

    /// Ensure the entry holds a [`Mapping`] and return it.
    ///
    /// If the entry is vacant, an empty mapping is inserted. If the entry is
    /// occupied but the value is not a mapping, it is replaced with an empty
    /// mapping. Otherwise the existing nested mapping is returned unchanged.
    pub fn or_insert_mapping(self) -> Mapping {
        match self {
            Entry::Occupied(o) => {
                if let Some(existing) = o.get_mapping() {
                    return existing;
                }
                o.set_value(Mapping::new());
                o.get_mapping()
                    .expect("value was just replaced with a mapping")
            }
            Entry::Vacant(v) => {
                let VacantEntry { mapping, key } = v;
                mapping.set(&key, Mapping::new());
                mapping
                    .get_mapping(&key)
                    .expect("mapping was just inserted")
            }
        }
    }

    /// Ensure the entry holds a [`Sequence`] and return it.
    ///
    /// If the entry is vacant, an empty sequence is inserted. If the entry is
    /// occupied but the value is not a sequence, it is replaced with an empty
    /// sequence. Otherwise the existing nested sequence is returned unchanged.
    pub fn or_insert_sequence(self) -> Sequence {
        match self {
            Entry::Occupied(o) => {
                if let Some(existing) = o.get_sequence() {
                    return existing;
                }
                o.set_value(Sequence::new());
                o.get_sequence()
                    .expect("value was just replaced with a sequence")
            }
            Entry::Vacant(v) => {
                let VacantEntry { mapping, key } = v;
                mapping.set(&key, Sequence::new());
                mapping
                    .get_sequence(&key)
                    .expect("sequence was just inserted")
            }
        }
    }

    /// Provide in-place access to an occupied entry before any potential
    /// inserts.
    ///
    /// The closure is only called if the entry is occupied. Returns `self`
    /// unchanged so calls can be chained with [`or_insert`](Self::or_insert)
    /// and friends.
    ///
    /// The tree uses interior mutability, so the closure receives a shared
    /// reference; `OccupiedEntry::set_value` and friends still work through
    /// `&self`.
    pub fn and_modify<F>(self, f: F) -> Self
    where
        F: FnOnce(&OccupiedEntry<'a>),
    {
        if let Entry::Occupied(ref o) = self {
            f(o);
        }
        self
    }

    /// Return the key that this entry refers to, as a [`YamlNode`].
    ///
    /// For an occupied entry this is the existing key node in the tree. For
    /// a vacant entry it is the key that would be inserted, rendered as a
    /// fresh scalar node.
    pub fn key(&self) -> Option<YamlNode> {
        match self {
            Entry::Occupied(o) => o.key(),
            Entry::Vacant(v) => key_as_yaml_node(&v.key),
        }
    }
}

impl<'a> OccupiedEntry<'a> {
    /// Return the key of this entry as a [`YamlNode`].
    ///
    /// Returns `None` for malformed entries with no key node.
    pub fn key(&self) -> Option<YamlNode> {
        self.entry.key_node()
    }

    /// Return the current value of this entry as a [`YamlNode`].
    ///
    /// Returns `None` for malformed entries with no value node.
    pub fn get(&self) -> Option<YamlNode> {
        self.entry.value_node()
    }

    /// Return the value as a nested [`Mapping`], if it is one, looking
    /// through a tag on it.
    pub fn get_mapping(&self) -> Option<Mapping> {
        self.get()?.mapping_content()
    }

    /// Return the value as a nested [`Sequence`], if it is one, looking
    /// through a tag on it.
    pub fn get_sequence(&self) -> Option<Sequence> {
        self.get()?.sequence_content()
    }

    /// Replace the value of this entry with `value`.
    pub fn set_value(&self, value: impl AsYaml) {
        let flow_context = self.mapping.is_flow_style();
        self.entry.set_value(value, flow_context);
    }

    /// Remove this entry from the mapping and return the removed
    /// [`MappingEntry`] handle. The returned entry is detached from the tree.
    pub fn remove_entry(self) -> MappingEntry {
        self.entry.clone().remove();
        self.entry
    }

    /// Consume the [`OccupiedEntry`] and return the current value as a
    /// [`YamlNode`].
    ///
    /// Panics if the entry has no value node (indicates malformed input).
    fn into_value(self) -> YamlNode {
        self.entry
            .value_node()
            .expect("occupied entry has no value node")
    }
}

impl<'a, K: AsYaml> VacantEntry<'a, K> {
    /// Return a reference to the pending key.
    pub fn key(&self) -> &K {
        &self.key
    }

    /// Consume the [`VacantEntry`] and return the key that would have been
    /// inserted.
    pub fn into_key(self) -> K {
        self.key
    }

    /// Insert `value` into the mapping at this vacant entry's key, and
    /// return the newly inserted value as a [`YamlNode`].
    pub fn insert(self, value: impl AsYaml) -> YamlNode {
        self.mapping.set(&self.key, value);
        self.mapping
            .get(&self.key)
            .expect("value just inserted must be retrievable")
    }
}

fn key_as_yaml_node<K: AsYaml>(key: &K) -> Option<YamlNode> {
    use crate::lex::SyntaxKind;
    use rowan::GreenNodeBuilder;

    if let Some(node) = key.value_node() {
        return YamlNode::from_syntax(node);
    }
    let mut builder = GreenNodeBuilder::new();
    builder.start_node(SyntaxKind::ROOT.into());
    key.build_content(&mut builder, 0, false);
    builder.finish_node();
    let root = crate::nodes::SyntaxNode::new_root(builder.finish());
    crate::nodes::held_node(&root)
}