fancy colours

4 files changed, 207 insertions, 1 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 7d6702a..5874327 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -857,9 +857,11 @@ version = "0.1.0"
 dependencies = [
  "clap",
  "miette",
+ "owo-colors",
  "protobuf",
  "protobuf-codegen",
  "protobuf-json-mapping",
+ "protobuf-support",
  "reqwest",
  "tokio",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 8fdd38e..3d7507c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,8 +11,10 @@ protobuf-codegen = "3.2"
 [dependencies]
 protobuf = "3.2"
 protobuf-json-mapping = "3.2"
+protobuf-support = "3.2"
 
 reqwest = { version = "0.11", features = ["json", "rustls-tls"], default-features = false }
 tokio = { version = "1", features = ["full"] }
 clap = { version = "4", features = ["derive"] }
 miette = { version = "5", features = ["fancy"] }
+owo-colors = "3.5"
diff --git a/src/fancy.rs b/src/fancy.rs
new file mode 100644
index 0000000..477018f
--- /dev/null
+++ b/src/fancy.rs
@@ -0,0 +1,200 @@
+
+use owo_colors::Style;
+use protobuf::reflect::*;
+use protobuf::MessageDyn;
+use protobuf::UnknownValueRef;
+use std::fmt;
+use std::fmt::Display;
+use std::fmt::Write;
+use owo_colors::OwoColorize;
+
+use protobuf_support::text_format::quote_bytes_to;
+
+
+pub fn print_to_string_fancy(m: &dyn MessageDyn) -> String {
+    let mut r = String::new();
+    print_to_internal(&MessageRef::from(m), &mut r, true, 0);
+    r.to_string()
+}
+
+fn print_str_to(s: &str, buf: &mut String) {
+    // TODO: keep printable Unicode
+    quote_bytes_to(s.as_bytes(), buf);
+}
+
+fn do_indent(buf: &mut String, pretty: bool, indent: usize) {
+    if pretty && indent > 0 {
+        for _ in 0..indent {
+            buf.push_str("  ");
+        }
+    }
+}
+
+trait FieldName: fmt::Display {}
+impl<'a> FieldName for &'a str {}
+impl FieldName for u32 {}
+
+fn print_start_field<F: FieldName>(
+    buf: &mut String,
+    pretty: bool,
+    indent: usize,
+    first: &mut bool,
+    field_name: F,
+) {
+    if !*first && !pretty {
+        buf.push_str(" ");
+    }
+    do_indent(buf, pretty, indent);
+    *first = false;
+    if indent == 0 {
+        write!(buf, "{}", field_name.bold().bright_red()).unwrap();
+    } else {
+        write!(buf, "{}", field_name.bright_blue()).unwrap();
+    }
+}
+
+fn print_end_field(buf: &mut String, pretty: bool) {
+    if pretty {
+        buf.push_str("\n");
+    }
+}
+
+
+struct Prefix (Style);
+struct Suffix (Style);
+
+impl Display for Prefix {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        self.0.fmt_prefix(fmt)
+    }
+}
+impl Display for Suffix {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        self.0.fmt_suffix(fmt)
+    }
+}
+
+fn print_field<F: FieldName>(
+    buf: &mut String,
+    pretty: bool,
+    indent: usize,
+    first: &mut bool,
+    field_name: F,
+    value: ReflectValueRef,
+) {
+    print_start_field(buf, pretty, indent, first, field_name);
+
+    match value {
+        ReflectValueRef::Message(m) => {
+            buf.push_str(" {");
+            if pretty {
+                buf.push_str("\n");
+            }
+            print_to_internal(&m, buf, pretty, indent + 1);
+            do_indent(buf, pretty, indent);
+            buf.push_str("}");
+        }
+        ReflectValueRef::Enum(d, v) => {
+            buf.push_str(": ");
+            match d.value_by_number(v) {
+                Some(e) => write!(buf, "{}", e.name().bright_yellow()).unwrap(),
+                None => write!(buf, ": {}", v.bright_red()).unwrap(),
+            }
+        }
+        ReflectValueRef::String(s) => {
+            let style = Style::new().bright_cyan();
+            buf.push_str(": ");
+            write!(buf, "{}", Prefix(style)).unwrap();
+            print_str_to(s, buf);
+            write!(buf, "{}", Suffix(style)).unwrap();
+        }
+        ReflectValueRef::Bytes(b) => {
+            buf.push_str(": ");
+            quote_bytes_to(b, buf);
+        }
+        ReflectValueRef::I32(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::I64(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::U32(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::U64(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::Bool(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::F32(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+        ReflectValueRef::F64(v) => {
+            write!(buf, ": {}", v.bright_yellow()).unwrap();
+        }
+    }
+
+    print_end_field(buf, pretty);
+}
+
+
+pub fn print_to_internal(m: &MessageRef, buf: &mut String, pretty: bool, indent: usize) {
+    let d = m.descriptor_dyn();
+    let mut first = true;
+    for f in d.fields() {
+        match f.get_reflect(&**m) {
+            ReflectFieldRef::Map(map) => {
+                for (k, v) in &map {
+                    print_start_field(buf, pretty, indent, &mut first, f.name());
+                    buf.push_str(" {");
+                    if pretty {
+                        buf.push_str("\n");
+                    }
+
+                    let mut entry_first = true;
+
+                    print_field(buf, pretty, indent + 1, &mut entry_first, "key", k);
+                    print_field(buf, pretty, indent + 1, &mut entry_first, "value", v);
+                    do_indent(buf, pretty, indent);
+                    buf.push_str("}");
+                    print_end_field(buf, pretty);
+                }
+            }
+            ReflectFieldRef::Repeated(repeated) => {
+                for v in repeated {
+                    print_field(buf, pretty, indent, &mut first, f.name(), v);
+                }
+            }
+            ReflectFieldRef::Optional(optional) => {
+                if let Some(v) = optional.value() {
+                    print_field(buf, pretty, indent, &mut first, f.name(), v);
+                }
+            }
+        }
+    }
+
+    let mut fields: Vec<(u32, UnknownValueRef)> = m.unknown_fields_dyn().iter().collect();
+    // Sort for stable output
+    fields.sort_by_key(|(field_number, _)| *field_number);
+    for (field_number, value) in fields {
+        // TODO: try decode nested message for length-delimited
+        print_field(
+            buf,
+            pretty,
+            indent,
+            &mut first,
+            field_number,
+            to_reflect_value_ref(&value),
+        );
+    }
+}
+
+fn to_reflect_value_ref<'o>(r: &'o UnknownValueRef<'o>) -> ReflectValueRef<'o> {
+    match r {
+        UnknownValueRef::Fixed32(v) => ReflectValueRef::U32(*v),
+        UnknownValueRef::Fixed64(v) => ReflectValueRef::U64(*v),
+        UnknownValueRef::Varint(v) => ReflectValueRef::U64(*v),
+        UnknownValueRef::LengthDelimited(v) => ReflectValueRef::Bytes(v),
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index feb6e9c..59e7e51 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,6 @@
 
 mod protos;
+mod fancy;
 use protos::protos::gtfs_realtime::FeedMessage;
 
 use protobuf::Message;
@@ -52,8 +53,9 @@ async fn main() -> miette::Result<()> {
         true => 
             println!("{}", protobuf_json_mapping::print_to_string(&proto).into_diagnostic()?),
         false => 
-            println!("{}", protobuf::text_format::print_to_string_pretty(&proto))
+            println!("{}", fancy::print_to_string_fancy(&proto))
     }
+
     Ok(())
 }