Move shared_traits to ruff_formatter (#4632)

crates/ruff_python_formatter/src/cst/mod.rs

@@ -19,15 +19,15 @@ pub(crate) mod visitor;
 type Ident = String;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct Attributed<T> {
-    pub range: TextRange,
-    pub node: T,
-    pub trivia: Vec<Trivia>,
-    pub parentheses: Parenthesize,
+pub(crate) struct Attributed<T> {
+    pub(crate) range: TextRange,
+    pub(crate) node: T,
+    pub(crate) trivia: Vec<Trivia>,
+    pub(crate) parentheses: Parenthesize,
 }
 
 impl<T> Attributed<T> {
-    pub fn new(range: TextRange, node: T) -> Self {
+    pub(crate) fn new(range: TextRange, node: T) -> Self {
         Self {
             range,
             node,
@@ -36,23 +36,19 @@ impl<T> Attributed<T> {
         }
     }
 
-    pub const fn range(&self) -> TextRange {
+    pub(crate) const fn range(&self) -> TextRange {
         self.range
     }
 
-    pub const fn start(&self) -> TextSize {
+    pub(crate) const fn start(&self) -> TextSize {
         self.range.start()
     }
 
-    pub const fn end(&self) -> TextSize {
+    pub(crate) const fn end(&self) -> TextSize {
         self.range.end()
     }
 
-    pub fn add_trivia(&mut self, trivia: Trivia) {
-        self.trivia.push(trivia);
-    }
-
-    pub fn id(&self) -> usize {
+    pub(crate) fn id(&self) -> usize {
         std::ptr::addr_of!(self.node) as usize
     }
 }
@@ -65,7 +61,7 @@ impl<T> Deref for Attributed<T> {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum ExprContext {
+pub(crate) enum ExprContext {
     Load,
     Store,
     Del,
@@ -82,7 +78,7 @@ impl From<ast::ExprContext> for ExprContext {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum BoolOpKind {
+pub(crate) enum BoolOpKind {
     And,
     Or,
 }
@@ -99,7 +95,7 @@ impl From<&ast::Boolop> for BoolOpKind {
 pub(crate) type BoolOp = Attributed<BoolOpKind>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum OperatorKind {
+pub(crate) enum OperatorKind {
     Add,
     Sub,
     Mult,
@@ -138,7 +134,7 @@ impl From<&ast::Operator> for OperatorKind {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum UnaryOpKind {
+pub(crate) enum UnaryOpKind {
     Invert,
     Not,
     UAdd,
@@ -159,7 +155,7 @@ impl From<&ast::Unaryop> for UnaryOpKind {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum CmpOpKind {
+pub(crate) enum CmpOpKind {
     Eq,
     NotEq,
     Lt,
@@ -208,7 +204,7 @@ impl From<(Vec<ast::Stmt>, &Locator<'_>)> for Body {
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum StmtKind {
+pub(crate) enum StmtKind {
     FunctionDef {
         name: Ident,
         args: Box<Arguments>,
@@ -338,7 +334,7 @@ pub enum StmtKind {
 pub(crate) type Stmt = Attributed<StmtKind>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum ExprKind {
+pub(crate) enum ExprKind {
     BoolOp {
         ops: Vec<BoolOp>,
         values: Vec<Expr>,
@@ -456,15 +452,15 @@ pub enum ExprKind {
 pub(crate) type Expr = Attributed<ExprKind>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct Comprehension {
-    pub target: Expr,
-    pub iter: Expr,
-    pub ifs: Vec<Expr>,
-    pub is_async: usize,
+pub(crate) struct Comprehension {
+    pub(crate) target: Expr,
+    pub(crate) iter: Expr,
+    pub(crate) ifs: Vec<Expr>,
+    pub(crate) is_async: usize,
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum ExcepthandlerKind {
+pub(crate) enum ExcepthandlerKind {
     ExceptHandler {
         type_: Option<Box<Expr>>,
         name: Option<Ident>,
@@ -475,7 +471,7 @@ pub enum ExcepthandlerKind {
 pub(crate) type Excepthandler = Attributed<ExcepthandlerKind>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub enum SliceIndexKind {
+pub(crate) enum SliceIndexKind {
     /// The index slot exists, but is empty.
     Empty,
     /// The index slot contains an expression.
@@ -485,57 +481,57 @@ pub enum SliceIndexKind {
 pub(crate) type SliceIndex = Attributed<SliceIndexKind>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct Arguments {
-    pub posonlyargs: Vec<Arg>,
-    pub args: Vec<Arg>,
-    pub vararg: Option<Box<Arg>>,
-    pub kwonlyargs: Vec<Arg>,
-    pub kw_defaults: Vec<Expr>,
-    pub kwarg: Option<Box<Arg>>,
-    pub defaults: Vec<Expr>,
+pub(crate) struct Arguments {
+    pub(crate) posonlyargs: Vec<Arg>,
+    pub(crate) args: Vec<Arg>,
+    pub(crate) vararg: Option<Box<Arg>>,
+    pub(crate) kwonlyargs: Vec<Arg>,
+    pub(crate) kw_defaults: Vec<Expr>,
+    pub(crate) kwarg: Option<Box<Arg>>,
+    pub(crate) defaults: Vec<Expr>,
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct ArgData {
-    pub arg: Ident,
-    pub annotation: Option<Box<Expr>>,
-    pub type_comment: Option<String>,
+pub(crate) struct ArgData {
+    pub(crate) arg: Ident,
+    pub(crate) annotation: Option<Box<Expr>>,
+    pub(crate) type_comment: Option<String>,
 }
 
 pub(crate) type Arg = Attributed<ArgData>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct KeywordData {
-    pub arg: Option<Ident>,
-    pub value: Expr,
+pub(crate) struct KeywordData {
+    pub(crate) arg: Option<Ident>,
+    pub(crate) value: Expr,
 }
 
 pub(crate) type Keyword = Attributed<KeywordData>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct AliasData {
-    pub name: Ident,
-    pub asname: Option<Ident>,
+pub(crate) struct AliasData {
+    pub(crate) name: Ident,
+    pub(crate) asname: Option<Ident>,
 }
 
 pub(crate) type Alias = Attributed<AliasData>;
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct Withitem {
-    pub context_expr: Expr,
-    pub optional_vars: Option<Box<Expr>>,
+pub(crate) struct Withitem {
+    pub(crate) context_expr: Expr,
+    pub(crate) optional_vars: Option<Box<Expr>>,
 }
 
 #[derive(Clone, Debug, PartialEq)]
-pub struct MatchCase {
-    pub pattern: Pattern,
-    pub guard: Option<Box<Expr>>,
-    pub body: Body,
+pub(crate) struct MatchCase {
+    pub(crate) pattern: Pattern,
+    pub(crate) guard: Option<Box<Expr>>,
+    pub(crate) body: Body,
 }
 
 #[allow(clippy::enum_variant_names)]
 #[derive(Clone, Debug, PartialEq)]
-pub enum PatternKind {
+pub(crate) enum PatternKind {
     MatchValue {
         value: Box<Expr>,
     },

crates/ruff_python_formatter/src/format/alias.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Alias;
 use crate::format::comments::end_of_line_comments;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatAlias<'a> {
+pub(crate) struct FormatAlias<'a> {
     item: &'a Alias,
 }
 

crates/ruff_python_formatter/src/format/arg.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Arg;
 use crate::format::comments::end_of_line_comments;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatArg<'a> {
+pub(crate) struct FormatArg<'a> {
     item: &'a Arg,
 }
 

crates/ruff_python_formatter/src/format/arguments.rs

@@ -1,11 +1,9 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{format_args, write, Format};
 
-use crate::context::ASTFormatContext;
 use crate::cst::Arguments;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatArguments<'a> {
+pub(crate) struct FormatArguments<'a> {
     item: &'a Arguments,
 }
 

crates/ruff_python_formatter/src/format/bool_op.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{BoolOp, BoolOpKind};
 use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatBoolOp<'a> {
+pub(crate) struct FormatBoolOp<'a> {
     item: &'a BoolOp,
 }
 

crates/ruff_python_formatter/src/format/builders.rs

@@ -1,10 +1,8 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{write, Format};
 use ruff_text_size::TextRange;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{Body, Stmt};
-use crate::shared_traits::AsFormat;
 use crate::trivia::{Relationship, TriviaKind};
 
 #[derive(Copy, Clone)]

crates/ruff_python_formatter/src/format/cmp_op.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{CmpOp, CmpOpKind};
 use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatCmpOp<'a> {
+pub(crate) struct FormatCmpOp<'a> {
     item: &'a CmpOp,
 }
 

crates/ruff_python_formatter/src/format/comments.rs

@@ -1,7 +1,6 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{write, Format};
 
-use crate::context::ASTFormatContext;
 use crate::cst::Attributed;
 use crate::format::builders::literal;
 use crate::trivia::TriviaKind;

crates/ruff_python_formatter/src/format/comprehension.rs

@@ -1,11 +1,9 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Comprehension;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatComprehension<'a> {
+pub(crate) struct FormatComprehension<'a> {
     item: &'a Comprehension,
 }
 

crates/ruff_python_formatter/src/format/excepthandler.rs

@@ -1,13 +1,11 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{Excepthandler, ExcepthandlerKind};
 use crate::format::builders::block;
 use crate::format::comments::end_of_line_comments;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatExcepthandler<'a> {
+pub(crate) struct FormatExcepthandler<'a> {
     item: &'a Excepthandler,
 }
 

crates/ruff_python_formatter/src/format/expr.rs

@@ -2,10 +2,9 @@
 
 use rustpython_parser::ast::{Constant, ConversionFlag};
 
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{format_args, write};
 
-use crate::context::ASTFormatContext;
 use crate::cst::{
     Arguments, BoolOp, CmpOp, Comprehension, Expr, ExprKind, Keyword, Operator, OperatorKind,
     SliceIndex, SliceIndexKind, UnaryOp, UnaryOpKind,
@@ -15,10 +14,9 @@ use crate::format::comments::{dangling_comments, end_of_line_comments, leading_c
 use crate::format::helpers::{is_self_closing, is_simple_power, is_simple_slice};
 use crate::format::numbers::{complex_literal, float_literal, int_literal};
 use crate::format::strings::string_literal;
-use crate::shared_traits::AsFormat;
 use crate::trivia::{Parenthesize, TriviaKind};
 
-pub struct FormatExpr<'a> {
+pub(crate) struct FormatExpr<'a> {
     item: &'a Expr,
 }
 

crates/ruff_python_formatter/src/format/keyword.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Keyword;
 use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatKeyword<'a> {
+pub(crate) struct FormatKeyword<'a> {
     item: &'a Keyword,
 }
 

crates/ruff_python_formatter/src/format/match_case.rs

@@ -1,13 +1,11 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::MatchCase;
 use crate::format::builders::block;
 use crate::format::comments::{end_of_line_comments, leading_comments};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatMatchCase<'a> {
+pub(crate) struct FormatMatchCase<'a> {
     item: &'a MatchCase,
 }
 

crates/ruff_python_formatter/src/format/numbers.rs

@@ -1,10 +1,9 @@
 use std::ops::{Add, Sub};
 
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{write, Format};
 use ruff_text_size::{TextRange, TextSize};
 
-use crate::context::ASTFormatContext;
 use crate::format::builders::literal;
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]

crates/ruff_python_formatter/src/format/operator.rs

@@ -1,12 +1,10 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{Operator, OperatorKind};
 use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatOperator<'a> {
+pub(crate) struct FormatOperator<'a> {
     item: &'a Operator,
 }
 

crates/ruff_python_formatter/src/format/pattern.rs

@@ -1,13 +1,11 @@
 use rustpython_parser::ast::Constant;
 
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{Pattern, PatternKind};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatPattern<'a> {
+pub(crate) struct FormatPattern<'a> {
     item: &'a Pattern,
 }
 

crates/ruff_python_formatter/src/format/stmt.rs

@@ -1,9 +1,8 @@
 #![allow(unused_variables, clippy::too_many_arguments)]
 
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{format_args, write};
 
-use crate::context::ASTFormatContext;
 use crate::cst::{
     Alias, Arguments, Body, Excepthandler, Expr, ExprKind, Keyword, MatchCase, Operator, Stmt,
     StmtKind, Withitem,
@@ -11,7 +10,6 @@ use crate::cst::{
 use crate::format::builders::{block, join_names};
 use crate::format::comments::{end_of_line_comments, leading_comments, trailing_comments};
 use crate::format::helpers::is_self_closing;
-use crate::shared_traits::AsFormat;
 
 fn format_break(f: &mut Formatter<ASTFormatContext>, stmt: &Stmt) -> FormatResult<()> {
     write!(f, [text("break")])?;
@@ -752,7 +750,7 @@ fn format_with_(
     Ok(())
 }
 
-pub struct FormatStmt<'a> {
+pub(crate) struct FormatStmt<'a> {
     item: &'a Stmt,
 }
 

crates/ruff_python_formatter/src/format/strings.rs

@@ -1,11 +1,10 @@
 use rustpython_parser::{Mode, Tok};
 
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::{write, Format};
 use ruff_python_ast::str::{leading_quote, trailing_quote};
 use ruff_text_size::TextRange;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Expr;
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]

crates/ruff_python_formatter/src/format/unary_op.rs

@@ -1,11 +1,9 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::{UnaryOp, UnaryOpKind};
-use crate::shared_traits::AsFormat;
 
-pub struct FormatUnaryOp<'a> {
+pub(crate) struct FormatUnaryOp<'a> {
     item: &'a UnaryOp,
 }
 

crates/ruff_python_formatter/src/format/withitem.rs

@@ -1,11 +1,9 @@
-use ruff_formatter::prelude::*;
+use crate::prelude::*;
 use ruff_formatter::write;
 
-use crate::context::ASTFormatContext;
 use crate::cst::Withitem;
-use crate::shared_traits::AsFormat;
 
-pub struct FormatWithitem<'a> {
+pub(crate) struct FormatWithitem<'a> {
     item: &'a Withitem,
 }
 

crates/ruff_python_formatter/src/lib.rs

@@ -17,8 +17,10 @@ mod cst;
 mod format;
 mod newlines;
 mod parentheses;
-pub mod shared_traits;
-pub mod trivia;
+mod prelude;
+mod trivia;
+
+include!("../../ruff_formatter/shared_traits.rs");
 
 pub fn fmt(contents: &str) -> Result<Formatted<ASTFormatContext>> {
     // Create a reusable locator.
@@ -179,7 +181,7 @@ mod tests {
 
     #[test]
     fn string_processing() {
-        use ruff_formatter::prelude::*;
+        use crate::prelude::*;
         use ruff_formatter::{format, format_args, write};
 
         struct FormatString<'a>(&'a str);

crates/ruff_python_formatter/src/prelude.rs

@@ -0,0 +1,3 @@
+#[allow(unused_imports)]
+pub(crate) use crate::{ASTFormatContext, AsFormat, FormattedIterExt as _, IntoFormat};
+pub(crate) use ruff_formatter::prelude::*;

crates/ruff_python_formatter/src/shared_traits.rs

@@ -1,98 +0,0 @@
-#![allow(clippy::all)]
-
-/// Used to get an object that knows how to format this object.
-pub trait AsFormat<Context> {
-    type Format<'a>: ruff_formatter::Format<Context>
-    where
-        Self: 'a;
-
-    /// Returns an object that is able to format this object.
-    fn format(&self) -> Self::Format<'_>;
-}
-
-/// Implement [`AsFormat`] for references to types that implement [`AsFormat`].
-impl<T, C> AsFormat<C> for &T
-where
-    T: AsFormat<C>,
-{
-    type Format<'a> = T::Format<'a> where Self: 'a;
-
-    fn format(&self) -> Self::Format<'_> {
-        AsFormat::format(&**self)
-    }
-}
-
-/// Used to convert this object into an object that can be formatted.
-///
-/// The difference to [`AsFormat`] is that this trait takes ownership of `self`.
-pub trait IntoFormat<Context> {
-    type Format: ruff_formatter::Format<Context>;
-
-    fn into_format(self) -> Self::Format;
-}
-
-/// Implement [`IntoFormat`] for [`Option`] when `T` implements [`IntoFormat`]
-///
-/// Allows to call format on optional AST fields without having to unwrap the
-/// field first.
-impl<T, Context> IntoFormat<Context> for Option<T>
-where
-    T: IntoFormat<Context>,
-{
-    type Format = Option<T::Format>;
-
-    fn into_format(self) -> Self::Format {
-        self.map(IntoFormat::into_format)
-    }
-}
-
-/// Formatting specific [`Iterator`] extensions
-pub trait FormattedIterExt {
-    /// Converts every item to an object that knows how to format it.
-    fn formatted<Context>(self) -> FormattedIter<Self, Self::Item, Context>
-    where
-        Self: Iterator + Sized,
-        Self::Item: IntoFormat<Context>,
-    {
-        FormattedIter {
-            inner: self,
-            options: std::marker::PhantomData,
-        }
-    }
-}
-
-impl<I> FormattedIterExt for I where I: std::iter::Iterator {}
-
-pub struct FormattedIter<Iter, Item, Context>
-where
-    Iter: Iterator<Item = Item>,
-{
-    inner: Iter,
-    options: std::marker::PhantomData<Context>,
-}
-
-impl<Iter, Item, Context> std::iter::Iterator for FormattedIter<Iter, Item, Context>
-where
-    Iter: Iterator<Item = Item>,
-    Item: IntoFormat<Context>,
-{
-    type Item = Item::Format;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        Some(self.inner.next()?.into_format())
-    }
-}
-
-impl<Iter, Item, Context> std::iter::FusedIterator for FormattedIter<Iter, Item, Context>
-where
-    Iter: std::iter::FusedIterator<Item = Item>,
-    Item: IntoFormat<Context>,
-{
-}
-
-impl<Iter, Item, Context> std::iter::ExactSizeIterator for FormattedIter<Iter, Item, Context>
-where
-    Iter: Iterator<Item = Item> + std::iter::ExactSizeIterator,
-    Item: IntoFormat<Context>,
-{
-}

crates/ruff_python_formatter/src/trivia.rs

@@ -9,7 +9,7 @@ use crate::cst::{
 };
 
 #[derive(Clone, Copy, Debug)]
-pub enum Node<'a> {
+pub(crate) enum Node<'a> {
     Alias(&'a Alias),
     Arg(&'a Arg),
     Body(&'a Body),
@@ -27,7 +27,7 @@ pub enum Node<'a> {
 }
 
 impl Node<'_> {
-    pub fn id(&self) -> usize {
+    pub(crate) fn id(&self) -> usize {
         match self {
             Node::Alias(node) => node.id(),
             Node::Arg(node) => node.id(),
@@ -46,7 +46,7 @@ impl Node<'_> {
         }
     }
 
-    pub fn start(&self) -> TextSize {
+    pub(crate) fn start(&self) -> TextSize {
         match self {
             Node::Alias(node) => node.start(),
             Node::Arg(node) => node.start(),
@@ -65,7 +65,7 @@ impl Node<'_> {
         }
     }
 
-    pub fn end(&self) -> TextSize {
+    pub(crate) fn end(&self) -> TextSize {
         match self {
             Node::Alias(node) => node.end(),
             Node::Arg(node) => node.end(),
@@ -86,7 +86,7 @@ impl Node<'_> {
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum TriviaTokenKind {
+pub(crate) enum TriviaTokenKind {
     OwnLineComment,
     EndOfLineComment,
     MagicTrailingComma,
@@ -95,23 +95,23 @@ pub enum TriviaTokenKind {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub struct TriviaToken {
-    pub range: TextRange,
-    pub kind: TriviaTokenKind,
+pub(crate) struct TriviaToken {
+    pub(crate) range: TextRange,
+    pub(crate) kind: TriviaTokenKind,
 }
 
 impl TriviaToken {
-    pub const fn start(&self) -> TextSize {
+    pub(crate) const fn start(&self) -> TextSize {
         self.range.start()
     }
 
-    pub const fn end(&self) -> TextSize {
+    pub(crate) const fn end(&self) -> TextSize {
         self.range.end()
     }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, is_macro::Is)]
-pub enum TriviaKind {
+pub(crate) enum TriviaKind {
     /// A Comment that is separated by at least one line break from the
     /// preceding token.
     ///
@@ -140,14 +140,14 @@ pub enum TriviaKind {
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, is_macro::Is)]
-pub enum Relationship {
+pub(crate) enum Relationship {
     Leading,
     Trailing,
     Dangling,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, is_macro::Is)]
-pub enum Parenthesize {
+pub(crate) enum Parenthesize {
     /// Always parenthesize the statement or expression.
     Always,
     /// Never parenthesize the statement or expression.
@@ -157,13 +157,13 @@ pub enum Parenthesize {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub struct Trivia {
-    pub kind: TriviaKind,
-    pub relationship: Relationship,
+pub(crate) struct Trivia {
+    pub(crate) kind: TriviaKind,
+    pub(crate) relationship: Relationship,
 }
 
 impl Trivia {
-    pub fn from_token(token: &TriviaToken, relationship: Relationship) -> Self {
+    pub(crate) fn from_token(token: &TriviaToken, relationship: Relationship) -> Self {
         match token.kind {
             TriviaTokenKind::MagicTrailingComma => Self {
                 kind: TriviaKind::MagicTrailingComma,
@@ -189,7 +189,7 @@ impl Trivia {
     }
 }
 
-pub fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec<TriviaToken> {
+pub(crate) fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec<TriviaToken> {
     let mut tokens = vec![];
     let mut prev_tok: Option<(&Tok, TextRange)> = None;
     let mut prev_semantic_tok: Option<(&Tok, TextRange)> = None;
@@ -731,14 +731,14 @@ fn sorted_child_nodes_inner<'a>(node: Node<'a>, result: &mut Vec<Node<'a>>) {
     }
 }
 
-pub fn sorted_child_nodes(node: Node) -> Vec<Node> {
+pub(crate) fn sorted_child_nodes(node: Node) -> Vec<Node> {
     let mut result = Vec::new();
     sorted_child_nodes_inner(node, &mut result);
 
     result
 }
 
-pub fn decorate_token<'a>(
+pub(crate) fn decorate_token<'a>(
     token: &TriviaToken,
     node: Node<'a>,
     enclosing_node: Option<Node<'a>>,
@@ -818,20 +818,20 @@ pub fn decorate_token<'a>(
 }
 
 #[derive(Debug, Default)]
-pub struct TriviaIndex {
-    pub alias: FxHashMap<usize, Vec<Trivia>>,
-    pub arg: FxHashMap<usize, Vec<Trivia>>,
-    pub body: FxHashMap<usize, Vec<Trivia>>,
-    pub bool_op: FxHashMap<usize, Vec<Trivia>>,
-    pub cmp_op: FxHashMap<usize, Vec<Trivia>>,
-    pub excepthandler: FxHashMap<usize, Vec<Trivia>>,
-    pub expr: FxHashMap<usize, Vec<Trivia>>,
-    pub keyword: FxHashMap<usize, Vec<Trivia>>,
-    pub operator: FxHashMap<usize, Vec<Trivia>>,
-    pub pattern: FxHashMap<usize, Vec<Trivia>>,
-    pub slice_index: FxHashMap<usize, Vec<Trivia>>,
-    pub stmt: FxHashMap<usize, Vec<Trivia>>,
-    pub unary_op: FxHashMap<usize, Vec<Trivia>>,
+pub(crate) struct TriviaIndex {
+    pub(crate) alias: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) arg: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) body: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) bool_op: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) cmp_op: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) excepthandler: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) expr: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) keyword: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) operator: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) pattern: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) slice_index: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) stmt: FxHashMap<usize, Vec<Trivia>>,
+    pub(crate) unary_op: FxHashMap<usize, Vec<Trivia>>,
 }
 
 fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
@@ -931,7 +931,7 @@ fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
     }
 }
 
-pub fn decorate_trivia(tokens: Vec<TriviaToken>, python_ast: &[Stmt]) -> TriviaIndex {
+pub(crate) fn decorate_trivia(tokens: Vec<TriviaToken>, python_ast: &[Stmt]) -> TriviaIndex {
     let mut stack = vec![];
     let mut cache = FxHashMap::default();
     for token in &tokens {