1 files changed, 23 insertions, 7 deletions

diff --git a/src/main.rs b/src/main.rs
index 74db5cd..22ea865 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,8 @@
-use std::{path::PathBuf, str::FromStr};
+use std::{
+    io::{stdout, BufWriter, Write},
+    path::PathBuf,
+    str::FromStr,
+};
 
 use syntect::{
     dumps::{dump_to_uncompressed_file, from_uncompressed_dump_file},
@@ -158,7 +162,9 @@ fn highlight() {
         .or_else(|| set.find_syntax_by_first_line(&line))
         .unwrap_or_else(|| set.find_syntax_plain_text());
 
-    print!(r#"<pre class="highlight"><code>"#);
+    let mut stdout = BufWriter::new(stdout().lock());
+
+    write!(&mut stdout, r#"<pre class="highlight"><code>"#).unwrap();
 
     let mut parse = ParseState::new(syntax);
     let mut stack = ScopeStack::new();
@@ -191,17 +197,27 @@ fn highlight() {
 
         for (section, class) in spans {
             if let Some(idx) = class {
-                print!(r#"<span class="hl-style{idx}">"#)
+                write!(&mut stdout, r#"<span class="hl-style{idx}">"#).unwrap();
+            }
+            for ch in section.chars() {
+                match ch {
+                    '>' => write!(&mut stdout, "&gt;").unwrap(),
+                    '<' => write!(&mut stdout, "&lt;").unwrap(),
+                    '&' => write!(&mut stdout, "&amp;").unwrap(),
+                    '\'' => write!(&mut stdout, "&#39;").unwrap(),
+                    '"' => write!(&mut stdout, "&quot;").unwrap(),
+                    _ => write!(&mut stdout, "{ch}").unwrap(),
+                }
             }
-            print!("{section}");
             if class.is_some() {
-                print!("</span>");
+                write!(&mut stdout, "</span>").unwrap();
             }
         }
 
+        stdout.flush().unwrap();
         line.clear();
         if std::io::stdin().read_line(&mut line).unwrap() == 0 {
-            println!();
+            write!(&mut stdout, "\n").unwrap();
             break;
         }
         if !line.ends_with('\n') {
@@ -209,5 +225,5 @@ fn highlight() {
         }
     }
 
-    print!("</code></pre>");
+    write!(&mut stdout, "</code></pre>").unwrap();
 }