apollo_federation/connectors/json_selection/

pretty.rs

//! Pretty printing utility methods
//!
//! Working with raw JSONSelections when doing snapshot testing is difficult to
//! read and makes the snapshots themselves quite large. This module adds a new
//! pretty printing trait which is then implemented on the various sub types
//! of the JSONSelection tree.

use itertools::Itertools;

use super::lit_expr::LitExpr;
use super::lit_expr::LitOp;
use super::location::WithRange;
use super::parser::Alias;
use super::parser::Key;
use crate::connectors::json_selection::JSONSelection;
use crate::connectors::json_selection::MethodArgs;
use crate::connectors::json_selection::NamedSelection;
use crate::connectors::json_selection::NamingPrefix;
use crate::connectors::json_selection::PathList;
use crate::connectors::json_selection::PathSelection;
use crate::connectors::json_selection::SubSelection;
use crate::connectors::json_selection::TopLevelSelection;

impl std::fmt::Display for JSONSelection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&PrettyPrintable::pretty_print(self))
    }
}

/// Pretty print trait
///
/// This trait marks a type as supporting pretty printing itself outside of a
/// Display implementation, which might be more useful for snapshots.
///
/// Output preserves what the user wrote: `PathList::Expr(inner, tail)` always
/// emits `$(inner)rest`, because a `PathList::Expr` node in the AST can only
/// come from an explicit `$(...)` wrapper in source (that invariant is
/// enforced by the parser). There is no spec-dependent rewriting.
pub(crate) trait PrettyPrintable {
    /// Pretty print the struct
    fn pretty_print(&self) -> String {
        self.pretty_print_with_indentation(false, 0)
    }

    /// Pretty print the struct, with indentation
    ///
    /// Each indentation level is marked with 2 spaces, with `inline` signifying
    /// that the first line should be not indented.
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String;
}

/// Helper method to generate indentation
fn indent_chars(indent: usize) -> String {
    "  ".repeat(indent)
}

impl PrettyPrintable for JSONSelection {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        match &self.inner {
            TopLevelSelection::Named(named) => named.print_subselections(inline, indentation),
            TopLevelSelection::Value(lit) => lit.pretty_print_with_indentation(inline, indentation),
        }
    }
}

impl PrettyPrintable for SubSelection {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let mut result = String::new();

        result.push('{');

        if self.selections.is_empty() {
            result.push('}');
            return result;
        }

        if inline {
            result.push(' ');
        } else {
            result.push('\n');
            result.push_str(indent_chars(indentation + 1).as_str());
        }

        result.push_str(&self.print_subselections(inline, indentation + 1));

        if inline {
            result.push(' ');
        } else {
            result.push('\n');
            result.push_str(indent_chars(indentation).as_str());
        }

        result.push('}');

        result
    }
}

impl SubSelection {
    /// Prints all of the selections in a subselection
    fn print_subselections(&self, inline: bool, indentation: usize) -> String {
        let separator = if inline {
            ' '.to_string()
        } else {
            format!("\n{}", indent_chars(indentation))
        };

        self.selections
            .iter()
            .map(|s| s.pretty_print_with_indentation(inline, indentation))
            .join(separator.as_str())
    }
}

impl PrettyPrintable for PathSelection {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let inner = self.path.pretty_print_with_indentation(inline, indentation);
        // Because we can't tell where PathList::Key elements appear in the path
        // once we're inside PathList::pretty_print_with_indentation, we print
        // all PathList::Key elements with a leading '.' character, but we
        // remove the initial '.' if the path has more than one element, because
        // then the leading '.' is not necessary to disambiguate the key from a
        // field. To complicate matters further, inner may begin with spaces due
        // to indentation.
        let leading_space_count = inner.chars().take_while(|c| *c == ' ').count();
        let suffix = inner[leading_space_count..].to_string();
        if let Some(after_dot) = suffix.strip_prefix('.') {
            // Strip the '.' but keep any leading spaces.
            format!("{}{}", " ".repeat(leading_space_count), after_dot)
        } else {
            inner
        }
    }
}

impl PrettyPrintable for PathList {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let mut result = String::new();

        match self {
            Self::Var(var, tail) => {
                let rest = tail.pretty_print_with_indentation(inline, indentation);
                result.push_str(var.as_str());
                result.push_str(rest.as_str());
            }
            Self::Key(key, tail) => {
                result.push('.');
                result.push_str(key.pretty_print().as_str());
                let rest = tail.pretty_print_with_indentation(inline, indentation);
                result.push_str(rest.as_str());
            }
            Self::Expr(expr, tail) => {
                // `PathList::Expr` comes only from an explicit `$(...)` in
                // source, so pretty-print it back verbatim — the wrapper is
                // part of what the user wrote, not a syntactic accident.
                let inner = expr.pretty_print_with_indentation(inline, indentation);
                let rest = tail.pretty_print_with_indentation(inline, indentation);
                result.push_str("$(");
                result.push_str(inner.as_str());
                result.push(')');
                result.push_str(rest.as_str());
            }
            Self::Method(method, args, tail) => {
                result.push_str("->");
                result.push_str(method.as_str());
                if let Some(args) = args {
                    result.push_str(
                        args.pretty_print_with_indentation(inline, indentation)
                            .as_str(),
                    );
                }
                result.push_str(
                    tail.pretty_print_with_indentation(inline, indentation)
                        .as_str(),
                );
            }
            Self::Question(tail) => {
                result.push('?');
                let rest = tail.pretty_print_with_indentation(true, indentation);
                result.push_str(rest.as_str());
            }
            Self::Selection(sub) => {
                let sub = sub.pretty_print_with_indentation(inline, indentation);
                result.push(' ');
                result.push_str(sub.as_str());
            }
            Self::Empty => {}
        }

        result
    }
}

impl PrettyPrintable for MethodArgs {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let printed_args: Vec<String> = self
            .args
            .iter()
            .map(|arg| arg.pretty_print_with_indentation(inline, indentation + 1))
            .collect();

        // Check if args would produce multi-line output in non-inline mode,
        // which determines whether we break across lines (or add spacing in
        // inline mode to be consistent with collapsed multi-line output).
        let would_break = if inline {
            self.args
                .iter()
                .any(|arg| arg.pretty_print_with_indentation(false, 0).contains('\n'))
        } else {
            printed_args.iter().any(|a| a.contains('\n'))
        };

        if !inline && would_break {
            let indent = indent_chars(indentation + 1);
            let separator = format!(",\n{indent}");
            let joined = printed_args.iter().map(String::as_str).join(&separator);
            format!("(\n{indent}{joined}\n{})", indent_chars(indentation))
        } else if would_break {
            let joined = printed_args.iter().map(String::as_str).join(", ");
            format!("( {joined} )")
        } else {
            let joined = printed_args.iter().map(String::as_str).join(", ");
            format!("({joined})")
        }
    }
}

impl LitExpr {
    fn is_shorthand_property(key: &WithRange<Key>, value: &WithRange<LitExpr>) -> bool {
        let Key::Field(key_name) = key.as_ref() else {
            return false;
        };
        let LitExpr::Path(PathSelection { path }) = value.as_ref() else {
            return false;
        };
        let PathList::Key(path_key, tail) = path.as_ref() else {
            return false;
        };
        let tail_is_simple = match tail.as_ref() {
            PathList::Empty => true,
            // Allow shorthand for optional paths like { a? } which desugar to { a: a? }
            PathList::Question(inner) => matches!(inner.as_ref(), PathList::Empty),
            // Allow shorthand for paths with subselections like { a { b c }, d }
            PathList::Selection(_) => true,
            _ => false,
        };
        tail_is_simple && path_key.as_str() == key_name
    }
}

impl PrettyPrintable for LitExpr {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let mut result = String::new();

        match self {
            Self::String(s) => {
                let safely_quoted = serde_json_bytes::Value::String(s.clone().into()).to_string();
                result.push_str(safely_quoted.as_str());
            }
            Self::Number(n) => result.push_str(n.to_string().as_str()),
            Self::Bool(b) => result.push_str(b.to_string().as_str()),
            Self::Null => result.push_str("null"),
            Self::Object(sub) => {
                result.push('{');

                if sub.selections.is_empty() {
                    result.push('}');
                    return result;
                }

                let mut is_first = true;
                for sel in &sub.selections {
                    if is_first {
                        is_first = false;
                    } else {
                        result.push(',');
                    }

                    if inline {
                        result.push(' ');
                    } else {
                        result.push('\n');
                        result.push_str(indent_chars(indentation + 1).as_str());
                    }

                    result.push_str(
                        sel.pretty_print_with_indentation(inline, indentation + 1)
                            .as_str(),
                    );
                }

                if inline {
                    result.push(' ');
                } else {
                    result.push('\n');
                    result.push_str(indent_chars(indentation).as_str());
                }

                result.push('}');
            }
            Self::LegacyObject(map) => {
                result.push('{');

                if map.is_empty() {
                    result.push('}');
                    return result;
                }

                let mut is_first = true;
                for (key, value) in map {
                    if is_first {
                        is_first = false;
                    } else {
                        result.push(',');
                    }

                    if inline {
                        result.push(' ');
                    } else {
                        result.push('\n');
                        result.push_str(indent_chars(indentation + 1).as_str());
                    }

                    if Self::is_shorthand_property(key, value) {
                        // Print just the value path, which includes any
                        // trailing `?` (e.g. `a?` for { a: a? }).
                        result.push_str(
                            value
                                .pretty_print_with_indentation(inline, indentation + 1)
                                .as_str(),
                        );
                    } else {
                        result.push_str(key.pretty_print().as_str());
                        result.push_str(": ");
                        result.push_str(
                            value
                                .pretty_print_with_indentation(inline, indentation + 1)
                                .as_str(),
                        );
                    }
                }

                if inline {
                    result.push(' ');
                } else {
                    result.push('\n');
                    result.push_str(indent_chars(indentation).as_str());
                }

                result.push('}');
            }
            Self::Array(vec) => {
                result.push('[');
                let mut is_first = true;
                for value in vec {
                    if is_first {
                        is_first = false;
                    } else {
                        result.push_str(", ");
                    }
                    result.push_str(
                        value
                            .pretty_print_with_indentation(inline, indentation)
                            .as_str(),
                    );
                }
                result.push(']');
            }
            Self::Path(path) => {
                result.push_str(
                    path.pretty_print_with_indentation(inline, indentation)
                        .as_str(),
                );
            }
            Self::LitPath(literal, subpath) => {
                result.push_str(
                    literal
                        .pretty_print_with_indentation(inline, indentation)
                        .as_str(),
                );
                result.push_str(
                    subpath
                        .pretty_print_with_indentation(inline, indentation)
                        .as_str(),
                );
            }
            Self::OpChain(op, operands) => {
                let op_str = match op.as_ref() {
                    LitOp::NullishCoalescing => " ?? ",
                    LitOp::NoneCoalescing => " ?! ",
                };

                for (i, operand) in operands.iter().enumerate() {
                    if i > 0 {
                        result.push_str(op_str);
                    }
                    result.push_str(
                        operand
                            .pretty_print_with_indentation(inline, indentation)
                            .as_str(),
                    );
                }
            }
        }

        result
    }
}

impl PrettyPrintable for NamedSelection {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let mut result = String::new();

        // Collapse `alias: key...` to `key...` when alias matches the path's
        // single leading key (e.g. `a: a?` prints as `a?`). Both forms are
        // semantically equivalent and round-trip through the parser.
        let is_shorthand = if let NamingPrefix::Alias(alias) = &self.prefix
            && let LitExpr::Path(path) = self.path.as_ref()
        {
            path.get_single_key()
                .is_some_and(|key| alias.name.as_ref() == key.as_ref())
        } else {
            false
        };

        match &self.prefix {
            NamingPrefix::None => {}
            NamingPrefix::Alias(alias) => {
                if !is_shorthand {
                    result.push_str(alias.pretty_print().as_str());
                    result.push(' ');
                }
            }
            NamingPrefix::Spread(token_range) => {
                if token_range.is_some() {
                    result.push_str("... ");
                }
            }
        };

        // The .trim_start() handles the case when self.path is just a
        // SubSelection (i.e., a NamedGroupSelection), since that PathList
        // variant typically prints a single leading space.
        let pretty_path = self.path.pretty_print_with_indentation(inline, indentation);
        result.push_str(pretty_path.trim_start());

        result
    }
}

impl PrettyPrintable for Alias {
    fn pretty_print_with_indentation(&self, inline: bool, indentation: usize) -> String {
        let mut result = String::new();

        let name = self.name.pretty_print_with_indentation(inline, indentation);
        result.push_str(name.as_str());
        result.push(':');

        result
    }
}

impl PrettyPrintable for Key {
    fn pretty_print_with_indentation(&self, _inline: bool, _indentation: usize) -> String {
        match self {
            Self::Field(name) => name.clone(),
            Self::Quoted(name) => serde_json_bytes::Value::String(name.as_str().into()).to_string(),
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::connectors::JSONSelection;
    use crate::connectors::PathSelection;
    use crate::connectors::SubSelection;
    use crate::connectors::json_selection::NamedSelection;
    use crate::connectors::json_selection::PrettyPrintable;
    use crate::connectors::json_selection::location::new_span;
    use crate::connectors::json_selection::location::new_span_with_spec;
    use crate::connectors::json_selection::pretty::indent_chars;
    use crate::connectors::spec::ConnectSpec;
    use crate::selection;

    // Test all valid pretty print permutations
    fn test_permutations(selection: impl PrettyPrintable, expected: &str) {
        let indentation = 4;
        let expected_indented = expected
            .lines()
            .map(|line| format!("{}{line}", indent_chars(indentation)))
            .collect::<Vec<_>>()
            .join("\n");
        let expected_indented = expected_indented.trim_start();

        let prettified = selection.pretty_print();
        assert_eq!(
            prettified, expected,
            "pretty printing did not match: {prettified} != {expected}"
        );

        let prettified_inline = selection.pretty_print_with_indentation(true, indentation);
        let expected_inline = collapse_spaces(expected);
        assert_eq!(
            prettified_inline.trim_start(),
            expected_inline.trim_start(),
            "pretty printing inline did not match: {prettified_inline} != {}",
            expected_indented.trim_start()
        );

        let prettified_indented = selection.pretty_print_with_indentation(false, indentation);
        assert_eq!(
            prettified_indented, expected_indented,
            "pretty printing indented did not match: {prettified_indented} != {expected_indented}"
        );
    }

    fn collapse_spaces(s: impl Into<String>) -> String {
        let pattern = regex::Regex::new(r"\s+").expect("valid regex");
        pattern.replace_all(s.into().as_str(), " ").to_string()
    }

    #[test]
    fn it_prints_a_named_selection() {
        let selections = [
            // Field
            "cool",
            "cool: beans",
            "cool: beans {\n  whoa\n}",
            // Path
            "cool: one.two.three",
            // Quoted
            r#"cool: "b e a n s""#,
            "cool: \"b e a n s\" {\n  a\n  b\n}",
            // Group
            "cool: {\n  a\n  b\n}",
        ];
        for selection in selections {
            let (unmatched, named_selection) = NamedSelection::parse(new_span(selection)).unwrap();
            assert!(
                unmatched.is_empty(),
                "static named selection was not fully parsed: '{selection}' ({named_selection:?}) had unmatched '{unmatched}'"
            );

            test_permutations(named_selection, selection);
        }
    }

    #[test]
    fn it_prints_a_path_selection() {
        // These round-trip assertions include `$(...)` wrappers around
        // literals/OpChains, which v0.2/v0.3 require to disambiguate from
        // field lookups. v0.4 drops the wrappers (see the v0.4 variant below).
        let paths = [
            // Var
            "$.one.two.three",
            "$this.a.b",
            "$this.id.first {\n  username\n}",
            // Key
            "$.first",
            "a.b.c.d.e",
            "one.two.three {\n  a\n  b\n}",
            "$.single {\n  x\n}",
            "results->slice($(-1)->mul($args.suffixLength))",
            "$(1234)->add($(5678)->mul(2))",
            "$(true)->and($(false)->not)",
            "$(12345678987654321)->div(111111111)->eq(111111111)",
            "$(\"Product\")->slice(0, $(4)->mul(-1))->eq(\"Pro\")",
            "$($args.unnecessary.parens)->eq(42)",
        ];
        for path in paths {
            let (unmatched, path_selection) =
                PathSelection::parse(new_span_with_spec(path, ConnectSpec::V0_2)).unwrap();
            assert!(
                unmatched.is_empty(),
                "static path was not fully parsed: '{path}' ({path_selection:?}) had unmatched '{unmatched}'"
            );

            test_permutations(path_selection, path);
        }
    }

    #[test]
    fn it_prints_a_sub_selection() {
        let sub = "{\n  a\n  b\n}";
        let (unmatched, sub_selection) = SubSelection::parse(new_span(sub)).unwrap();
        assert!(
            unmatched.is_empty(),
            "static path was not fully parsed: '{sub}' ({sub_selection:?}) had unmatched '{unmatched}'"
        );

        test_permutations(sub_selection, sub);
    }

    #[test]
    fn it_prints_an_inline_path_with_subselection() {
        // This test ensures we do not print a leading ... before some.path,
        // even though we will probably want to do so once ... and conditional
        // selections are implemented. The printing of the leading ... was due
        // to an incomplete removal of experimental support for conditional
        // selections, which was put on hold to de-risk the GA release.
        let source = "before\nsome.path {\n  inline\n  me\n}\nafter";
        let sel = JSONSelection::parse(source).unwrap();
        test_permutations(sel, source);
    }

    #[test]
    fn it_prints_a_nested_sub_selection() {
        let sub = "{
          a {
            b {
              c
            }
          }
        }";
        let sub_indented = "{\n  a {\n    b {\n      c\n    }\n  }\n}";
        let sub_super_indented = "        {\n          a {\n            b {\n              c\n            }\n          }\n        }";

        let (unmatched, sub_selection) = SubSelection::parse(new_span(sub)).unwrap();

        assert!(
            unmatched.is_empty(),
            "static nested sub was not fully parsed: '{sub}' ({sub_selection:?}) had unmatched '{unmatched}'"
        );

        let pretty = sub_selection.pretty_print();
        assert_eq!(
            pretty, sub_indented,
            "nested sub pretty printing did not match: {pretty} != {sub_indented}"
        );

        let pretty = sub_selection.pretty_print_with_indentation(false, 4);
        assert_eq!(
            pretty,
            sub_super_indented.trim_start(),
            "nested inline sub pretty printing did not match: {pretty} != {}",
            sub_super_indented.trim_start()
        );
    }

    #[test]
    fn it_prints_root_selection() {
        let root_selection = JSONSelection::parse("id name").unwrap();
        test_permutations(root_selection, "id\nname");
    }

    // Round-trip helpers for connect/v0.4 top-level LitExpr forms. The
    // pretty-printer preserves whatever `$(...)` the user wrote (or didn't),
    // so for these canonical forms parse+pretty is the identity.
    fn assert_round_trip_v0_4(input: &str, canonical: &str) {
        let first = selection!(input, ConnectSpec::V0_4);
        let printed = first.pretty_print();
        assert_eq!(
            printed, canonical,
            "V0_4 pretty-print of `{input}` gave `{printed}`, expected `{canonical}`",
        );
        // Re-parsing the canonical form must yield a structurally identical
        // selection, or the pretty/parse cycle is asymmetric.
        let second = selection!(canonical, ConnectSpec::V0_4);
        let second_printed = second.pretty_print();
        assert_eq!(
            printed, second_printed,
            "re-parse of canonical V0_4 form did not reproduce the same AST",
        );
    }

    // The top-level `$(...)` form preserves its wrapper through pretty-print:
    // `PathList::Expr` is reserved for actual `$(...)` in source, so parsing
    // then pretty-printing is the identity in every case below.

    #[test]
    fn top_level_v0_4_number_literal() {
        assert_round_trip_v0_4("$(1234)", "$(1234)");
    }

    #[test]
    fn top_level_v0_4_negative_number_literal() {
        assert_round_trip_v0_4("$(-1)", "$(-1)");
    }

    #[test]
    fn top_level_v0_4_string_literal() {
        assert_round_trip_v0_4(r#"$("hello")"#, r#"$("hello")"#);
    }

    #[test]
    fn top_level_v0_4_boolean_literal() {
        assert_round_trip_v0_4("$(true)", "$(true)");
    }

    #[test]
    fn top_level_v0_4_null_literal() {
        assert_round_trip_v0_4("$(null)", "$(null)");
    }

    #[test]
    fn top_level_v0_4_array_literal() {
        assert_round_trip_v0_4("$([1, 2, 3])", "$([1, 2, 3])");
    }

    #[test]
    fn top_level_v0_4_number_with_method_chain() {
        assert_round_trip_v0_4("$(1234)->add(1111)", "$(1234)->add(1111)");
    }

    #[test]
    fn top_level_v0_4_negative_number_with_method_chain() {
        assert_round_trip_v0_4("$(-1)->add(10)", "$(-1)->add(10)");
    }

    #[test]
    fn top_level_v0_4_string_with_method_chain() {
        assert_round_trip_v0_4(r#"$("abc")->first"#, r#"$("abc")->first"#);
    }

    #[test]
    fn top_level_v0_4_array_with_method_chain() {
        assert_round_trip_v0_4("$([1, 2, 3])->last", "$([1, 2, 3])->last");
    }

    #[test]
    fn top_level_v0_4_object_literal_with_key_access() {
        // The `$({ a: 1, b: 2 }.b)` form wraps a `LitPath(Object, Key)`
        // expression. The pretty-printer formats the object multi-line, so
        // the canonical reflects that.
        assert_round_trip_v0_4("$({ a: 1, b: 2 }.b)", "$({\n  a: 1,\n  b: 2\n}.b)");
    }

    #[test]
    fn top_level_v0_4_operator_chain() {
        assert_round_trip_v0_4(
            r#"$($args.maybe ?? "fallback")"#,
            r#"$($args.maybe ?? "fallback")"#,
        );
    }

    #[test]
    fn top_level_v0_4_nested_expr_path_preserves_wrappers() {
        // `$($(-1)->add($(10)))` has three `$(...)` layers in source, and the
        // pretty-printer preserves every one of them — we never synthesize or
        // elide `$(...)` at any depth.
        assert_round_trip_v0_4("$($(-1)->add($(10)))", "$($(-1)->add($(10)))");
    }

    #[test]
    fn top_level_v0_4_bare_single_key_still_named() {
        // Single-key ambiguity: a bare key at the top level must still
        // parse as a NamedSelectionList so `author` means
        // `{ author: $.author }`, not a flat value lookup.
        let sel = selection!("author", ConnectSpec::V0_4);
        assert_eq!(sel.pretty_print(), "author");
        assert!(
            sel.next_subselection().is_some(),
            "top-level `author` should remain a NamedSelectionList wrapping a SubSelection",
        );
    }

    #[test]
    fn it_reprints_shorthand_properties() {
        let expected = r#"
upc
... category->match(
  ["book", {
    __typename: "Book",
    title,
    author {
      id
    }
  }],
  ["film", $ {
    __typename: "Film"
    title
    director {
      id
    }
  }],
  [@, null]
)"#
        .trim_start();

        let sel = selection!(&expected, ConnectSpec::V0_4);
        crate::assert_debug_snapshot!(&sel);

        test_permutations(sel, expected);
    }
}