1 files changed, 26 insertions, 8 deletions

diff --git a/src/writer.rs b/src/writer.rs
index 013d37b..f622bcf 100644
--- a/src/writer.rs
+++ b/src/writer.rs
@@ -13,6 +13,7 @@ use tokio::io::{AsyncWrite, AsyncWriteExt};
 use crate::{
     declaration::{Declaration, VersionInfo},
     element::{escape_str, Content, Element, IntoContent, IntoElement, Name, NamespaceDeclaration},
+    endable::Endable,
     error::Error,
     xml::{self, composers::Composer, parsers_complete::Parser, ETag, XMLDecl},
     Result, XMLNS_NS, XML_NS,
@@ -21,7 +22,7 @@ use crate::{
 // pub struct Writer<W, C = Composer> {
 #[derive(Debug)]
 pub struct Writer<W> {
-    inner: W,
+    inner: Endable<W>,
     depth: Vec<Name>,
     namespace_declarations: Vec<HashSet<NamespaceDeclaration>>,
 }
@@ -38,19 +39,20 @@ impl<W> Writer<W> {
             namespace: XMLNS_NS.to_string(),
         });
         Self {
-            inner: writer,
+            inner: Endable::new(writer),
             depth: Vec::new(),
             namespace_declarations: vec![default_declarations],
         }
     }
 
     pub fn into_inner(self) -> W {
-        self.inner
+        self.inner.into_inner()
     }
 }
 
 impl<W: AsyncWrite + Unpin + Send> Writer<W> {
     pub async fn write_declaration(&mut self, version: VersionInfo) -> Result<()> {
+        let writer = self.inner.try_as_mut()?;
         let declaration = Declaration::version(version);
         let version_info;
         match declaration.version_info {
@@ -64,7 +66,7 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
             encoding_decl: None,
             sd_decl: None,
         };
-        declaration.write(&mut self.inner).await?;
+        declaration.write(writer).await?;
         Ok(())
     }
 
@@ -105,6 +107,7 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
     }
 
     pub async fn write_empty(&mut self, element: &Element) -> Result<()> {
+        let writer = self.inner.try_as_mut()?;
         let mut namespace_declarations_stack: Vec<_> = self
             .namespace_declarations
             .iter()
@@ -204,12 +207,17 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
 
         let tag = xml::EmptyElemTag { name, attributes };
 
-        tag.write(&mut self.inner).await?;
+        tag.write(writer).await?;
+
+        if self.depth.is_empty() {
+            self.inner.end();
+        }
 
         Ok(())
     }
 
     pub async fn write_element_start(&mut self, element: &Element) -> Result<()> {
+        let writer = self.inner.try_as_mut()?;
         let mut namespace_declarations_stack: Vec<_> = self
             .namespace_declarations
             .iter()
@@ -309,7 +317,7 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
 
         let s_tag = xml::STag { name, attributes };
 
-        s_tag.write(&mut self.inner).await?;
+        s_tag.write(writer).await?;
 
         self.depth.push(element.name.clone());
         self.namespace_declarations
@@ -320,7 +328,12 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
     pub async fn write_content(&mut self, content: &Content) -> Result<()> {
         match content {
             Content::Element(element) => self.write_element(element).await?,
-            Content::Text(text) => self.inner.write_all(escape_str(text).as_bytes()).await?,
+            Content::Text(text) => {
+                self.inner
+                    .try_as_mut()?
+                    .write_all(escape_str(text).as_bytes())
+                    .await?
+            }
             // TODO: comments and PI
             Content::PI => {}
             Content::Comment(_) => {}
@@ -329,6 +342,7 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
     }
 
     pub async fn write_end(&mut self) -> Result<()> {
+        let writer = self.inner.try_as_mut()?;
         if let Some(name) = &self.depth.pop() {
             let e_tag;
             let namespace_declarations_stack: Vec<_> =
@@ -359,8 +373,12 @@ impl<W: AsyncWrite + Unpin + Send> Writer<W> {
                     )?),
                 };
             }
-            e_tag.write(&mut self.inner).await?;
+            e_tag.write(writer).await?;
             self.namespace_declarations.pop();
+
+            if self.depth.is_empty() {
+                self.inner.end();
+            }
             Ok(())
         } else {
             return Err(Error::NotInElement("".to_string()));