1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use std::fmt;

use crate::{Directive, StringValue};

/// UnionDefinitions are an abstract type where no common fields are declared.
///
/// *UnionDefTypeDefinition*:
///     Description? **union** Name Directives? UnionDefMemberTypes?
///
/// Detailed documentation can be found in [GraphQL spec](https://spec.graphql.org/October2021/#UnionTypeDefinition).
///
/// ### Example
/// ```rust
/// use apollo_encoder::UnionDefinition;
///
/// let mut union_ = UnionDefinition::new("Pet".to_string());
/// union_.member("Cat".to_string());
/// union_.member("Dog".to_string());
///
/// assert_eq!(
///     union_.to_string(),
/// r#"union Pet = Cat | Dog
/// "#
/// );
/// ```
#[derive(Debug, PartialEq, Clone)]
pub struct UnionDefinition {
    // Name must return a String.
    name: String,
    // Description may return a String.
    description: Option<StringValue>,
    // The vector of members that can be represented within this union.
    members: Vec<String>,
    /// Contains all directives.
    directives: Vec<Directive>,
    extend: bool,
}

impl UnionDefinition {
    /// Create a new instance of a UnionDef.
    pub fn new(name: String) -> Self {
        Self {
            name,
            description: None,
            members: Vec::new(),
            extend: false,
            directives: Vec::new(),
        }
    }

    /// Set the union type as an extension
    pub fn extend(&mut self) {
        self.extend = true;
    }

    /// Set the UnionDefs description.
    pub fn description(&mut self, description: String) {
        self.description = Some(StringValue::Top {
            source: description,
        });
    }

    /// Add a directive
    pub fn directive(&mut self, directive: Directive) {
        self.directives.push(directive);
    }

    /// Set a UnionDef member.
    pub fn member(&mut self, member: String) {
        self.members.push(member);
    }
}

impl fmt::Display for UnionDefinition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.extend {
            write!(f, "extend ")?;
        // No description when it's a extension
        } else if let Some(description) = &self.description {
            writeln!(f, "{description}")?;
        }

        write!(f, "union {}", self.name)?;

        for directive in &self.directives {
            write!(f, " {directive}")?;
        }

        write!(f, " =")?;

        for (i, member) in self.members.iter().enumerate() {
            match i {
                0 => write!(f, " {member}")?,
                _ => write!(f, " | {member}")?,
            }
        }

        writeln!(f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn it_encodes_union_with_description() {
        let mut union_ = UnionDefinition::new("Pet".to_string());
        union_.description("A union of all animals in a household.".to_string());
        union_.member("Cat".to_string());
        union_.member("Dog".to_string());

        assert_eq!(
            union_.to_string(),
            r#""A union of all animals in a household."
union Pet = Cat | Dog
"#
        );
    }

    #[test]
    fn it_encodes_union_extension() {
        let mut union_ = UnionDefinition::new("Pet".to_string());
        union_.description("A union of all animals in a household.".to_string());
        union_.member("Cat".to_string());
        union_.member("Dog".to_string());
        union_.extend();

        assert_eq!(
            union_.to_string(),
            r#"extend union Pet = Cat | Dog
"#
        );
    }
}