From 0ef6af807b74cf38c77b76e7c139a2492cb4b824 Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Fri, 11 Aug 2023 12:41:48 +0200
Subject: [PATCH] Implement DerefMut for WithNodeLevel (#6443)

**Summary** Implement `DerefMut` for `WithNodeLevel` so it can be used
in the same way as `PyFormatter`. I want this for my WIP upstack branch
to enable `.fmt(f)` on `WithNodeLevel` context. We could extend this to
remove the other two method from `WithNodeLevel`.
---
 crates/ruff_python_formatter/src/context.rs | 28 ++++++++++++++-------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/crates/ruff_python_formatter/src/context.rs b/crates/ruff_python_formatter/src/context.rs
index 23a3530a24..01caed9e68 100644
--- a/crates/ruff_python_formatter/src/context.rs
+++ b/crates/ruff_python_formatter/src/context.rs
@@ -1,9 +1,9 @@
 use crate::comments::Comments;
 use crate::PyFormatOptions;
-use ruff_formatter::prelude::*;
-use ruff_formatter::{Arguments, Buffer, FormatContext, GroupId, SourceCode};
+use ruff_formatter::{Buffer, FormatContext, GroupId, SourceCode};
 use ruff_source_file::Locator;
 use std::fmt::{Debug, Formatter};
+use std::ops::{Deref, DerefMut};
 
 #[derive(Clone)]
 pub struct PyFormatContext<'a> {
@@ -96,6 +96,7 @@ impl NodeLevel {
     }
 }
 
+/// Change the [`NodeLevel`] of the formatter for the lifetime of this struct
 pub(crate) struct WithNodeLevel<'ast, 'buf, B>
 where
     B: Buffer<Context = PyFormatContext<'ast>>,
@@ -119,16 +120,25 @@ where
             saved_level,
         }
     }
+}
 
-    #[inline]
-    pub(crate) fn write_fmt(&mut self, arguments: Arguments<B::Context>) -> FormatResult<()> {
-        self.buffer.write_fmt(arguments)
+impl<'ast, 'buf, B> Deref for WithNodeLevel<'ast, 'buf, B>
+where
+    B: Buffer<Context = PyFormatContext<'ast>>,
+{
+    type Target = B;
+
+    fn deref(&self) -> &Self::Target {
+        self.buffer
     }
+}
 
-    #[allow(unused)]
-    #[inline]
-    pub(crate) fn write_element(&mut self, element: FormatElement) -> FormatResult<()> {
-        self.buffer.write_element(element)
+impl<'ast, 'buf, B> DerefMut for WithNodeLevel<'ast, 'buf, B>
+where
+    B: Buffer<Context = PyFormatContext<'ast>>,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.buffer
     }
 }