From 0ef2853f27ea84407789cbd0680956f9e3ee9168 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kat=20March=C3=A1n?= <kzm@zkat.tech>
Date: Sun, 19 Sep 2021 18:41:21 -0700
Subject: [PATCH] feat(labels): replace snippet stuff with simpler labels (#62)

---
 Cargo.toml                       |   4 +-
 miette-derive/src/diagnostic.rs  |  47 ++--
 miette-derive/src/forward.rs     |  14 +-
 miette-derive/src/label.rs       | 179 +++++++++++++++
 miette-derive/src/lib.rs         |   5 +-
 miette-derive/src/snippets.rs    | 375 -------------------------------
 miette-derive/src/source_code.rs |  81 +++++++
 src/eyreish/context.rs           |  14 +-
 src/eyreish/wrapper.rs           |  26 ++-
 src/handlers/debug.rs            |   6 +-
 src/handlers/graphical.rs        | 135 +++++------
 src/handlers/narratable.rs       |  80 +------
 src/named_source.rs              |  42 ++--
 src/protocol.rs                  | 145 ++++++++----
 src/source_impls.rs              | 194 ++++++++++++----
 tests/derive.rs                  | 121 ++--------
 tests/narrated.rs                |  15 +-
 17 files changed, 729 insertions(+), 754 deletions(-)
 create mode 100644 miette-derive/src/label.rs
 delete mode 100644 miette-derive/src/snippets.rs
 create mode 100644 miette-derive/src/source_code.rs

diff --git a/Cargo.toml b/Cargo.toml
index 0533533..3703f79 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ unicode-width = { version = "0.1.8", optional = true }
 supports-hyperlinks = { version = "1.1.0", optional = true }
 supports-color = { version = "1.0.2", optional = true }
 supports-unicode = { version = "1.0.0", optional = true }
+itertools = { version = "0.10.1", optional = true }
 
 [dev-dependencies]
 semver = "1.0.4"
@@ -47,7 +48,8 @@ fancy = [
     "unicode-width",
     "supports-hyperlinks",
     "supports-color",
-    "supports-unicode"
+    "supports-unicode",
+    "itertools"
 ]
 
 [workspace]
diff --git a/miette-derive/src/diagnostic.rs b/miette-derive/src/diagnostic.rs
index db75868..50d47c4 100644
--- a/miette-derive/src/diagnostic.rs
+++ b/miette-derive/src/diagnostic.rs
@@ -6,8 +6,9 @@ use crate::code::Code;
 use crate::diagnostic_arg::DiagnosticArg;
 use crate::forward::{Forward, WhichFn};
 use crate::help::Help;
+use crate::label::Labels;
 use crate::severity::Severity;
-use crate::snippets::Snippets;
+use crate::source_code::SourceCode;
 use crate::url::Url;
 
 pub enum Diagnostic {
@@ -32,7 +33,7 @@ pub struct DiagnosticDef {
 
 pub enum DiagnosticDefArgs {
     Transparent(Forward),
-    Concrete(DiagnosticConcreteArgs),
+    Concrete(Box<DiagnosticConcreteArgs>),
 }
 
 impl DiagnosticDefArgs {
@@ -59,7 +60,8 @@ pub struct DiagnosticConcreteArgs {
     pub code: Option<Code>,
     pub severity: Option<Severity>,
     pub help: Option<Help>,
-    pub snippets: Option<Snippets>,
+    pub labels: Option<Labels>,
+    pub source_code: Option<SourceCode>,
     pub url: Option<Url>,
     pub forward: Option<Forward>,
 }
@@ -99,14 +101,16 @@ impl DiagnosticConcreteArgs {
                 }
             }
         }
-        let snippets = Snippets::from_fields(fields)?;
+        let labels = Labels::from_fields(fields)?;
+        let source_code = SourceCode::from_fields(fields)?;
         let concrete = DiagnosticConcreteArgs {
             code,
             help,
             severity,
-            snippets,
+            labels,
             url,
             forward,
+            source_code,
         };
         Ok(concrete)
     }
@@ -141,7 +145,7 @@ impl DiagnosticDefArgs {
             .into_iter()
             .filter(|x| !matches!(x, DiagnosticArg::Transparent));
         let concrete = DiagnosticConcreteArgs::parse(ident, fields, attr, args)?;
-        Ok(DiagnosticDefArgs::Concrete(concrete))
+        Ok(DiagnosticDefArgs::Concrete(Box::new(concrete)))
     }
 }
 
@@ -208,16 +212,18 @@ impl Diagnostic {
                         let code_method = forward.gen_struct_method(WhichFn::Code);
                         let help_method = forward.gen_struct_method(WhichFn::Help);
                         let url_method = forward.gen_struct_method(WhichFn::Url);
+                        let labels_method = forward.gen_struct_method(WhichFn::Labels);
+                        let source_code_method = forward.gen_struct_method(WhichFn::SourceCode);
                         let severity_method = forward.gen_struct_method(WhichFn::Severity);
-                        let snippets_method = forward.gen_struct_method(WhichFn::Snippets);
 
                         quote! {
                             impl #impl_generics miette::Diagnostic for #ident #ty_generics #where_clause {
                                 #code_method
                                 #help_method
                                 #url_method
+                                #labels_method
                                 #severity_method
-                                #snippets_method
+                                #source_code_method
                             }
                         }
                     }
@@ -243,24 +249,29 @@ impl Diagnostic {
                             .as_ref()
                             .and_then(|x| x.gen_struct())
                             .or_else(|| forward(WhichFn::Severity));
-                        let snip_body = concrete
-                            .snippets
-                            .as_ref()
-                            .and_then(|x| x.gen_struct(fields))
-                            .or_else(|| forward(WhichFn::Snippets));
                         let url_body = concrete
                             .url
                             .as_ref()
                             .and_then(|x| x.gen_struct(ident, fields))
                             .or_else(|| forward(WhichFn::Url));
-
+                        let labels_body = concrete
+                            .labels
+                            .as_ref()
+                            .and_then(|x| x.gen_struct(fields))
+                            .or_else(|| forward(WhichFn::Labels));
+                        let src_body = concrete
+                            .source_code
+                            .as_ref()
+                            .and_then(|x| x.gen_struct(fields))
+                            .or_else(|| forward(WhichFn::SourceCode));
                         quote! {
                             impl #impl_generics miette::Diagnostic for #ident #ty_generics #where_clause {
                                 #code_body
                                 #help_body
                                 #sev_body
-                                #snip_body
                                 #url_body
+                                #labels_body
+                                #src_body
                             }
                         }
                     }
@@ -275,14 +286,16 @@ impl Diagnostic {
                 let code_body = Code::gen_enum(variants);
                 let help_body = Help::gen_enum(variants);
                 let sev_body = Severity::gen_enum(variants);
-                let snip_body = Snippets::gen_enum(variants);
+                let labels_body = Labels::gen_enum(variants);
+                let src_body = SourceCode::gen_enum(variants);
                 let url_body = Url::gen_enum(ident, variants);
                 quote! {
                     impl #impl_generics miette::Diagnostic for #ident #ty_generics #where_clause {
                         #code_body
                         #help_body
                         #sev_body
-                        #snip_body
+                        #labels_body
+                        #src_body
                         #url_body
                     }
                 }
diff --git a/miette-derive/src/forward.rs b/miette-derive/src/forward.rs
index 36177d6..3d765ab 100644
--- a/miette-derive/src/forward.rs
+++ b/miette-derive/src/forward.rs
@@ -35,7 +35,8 @@ pub enum WhichFn {
     Help,
     Url,
     Severity,
-    Snippets,
+    Labels,
+    SourceCode,
 }
 
 impl WhichFn {
@@ -45,7 +46,8 @@ impl WhichFn {
             Self::Help => quote! { help() },
             Self::Url => quote! { url() },
             Self::Severity => quote! { severity() },
-            Self::Snippets => quote! { snippets() },
+            Self::Labels => quote! { labels() },
+            Self::SourceCode => quote! { source_code() },
         }
     }
 
@@ -63,11 +65,11 @@ impl WhichFn {
             Self::Severity => quote! {
                 fn severity(&self) -> std::option::Option<miette::Severity>
             },
-            Self::Snippets => quote! {
-                fn snippets(&self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = miette::DiagnosticSnippet> + '_>>
+            Self::Labels => quote! {
+                fn labels(&self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = miette::LabeledSpan> + '_>>
             },
-            Self::Related => quote! {
-                fn related<'a>(&'a self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = &'a dyn miette::Diagnostic> + 'a>>
+            Self::SourceCode => quote! {
+                fn source_code(&self) -> std::option::Option<&dyn miette::SourceCode>
             },
         }
     }
diff --git a/miette-derive/src/label.rs b/miette-derive/src/label.rs
new file mode 100644
index 0000000..5aed762
--- /dev/null
+++ b/miette-derive/src/label.rs
@@ -0,0 +1,179 @@
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote};
+use syn::{
+    parenthesized,
+    parse::{Parse, ParseStream},
+    spanned::Spanned,
+    Token,
+};
+
+use crate::{
+    diagnostic::{DiagnosticConcreteArgs, DiagnosticDef},
+    fmt::{self, Display},
+    forward::WhichFn,
+    utils::{display_pat_members, gen_all_variants_with},
+};
+
+pub struct Labels(Vec<Label>);
+
+struct Label {
+    label: Option<Display>,
+    span: syn::Member,
+}
+
+struct LabelAttr {
+    label: Option<Display>,
+}
+
+impl Parse for LabelAttr {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let la = input.lookahead1();
+        let label = if la.peek(syn::token::Paren) {
+            // #[label("{}", x)]
+            let content;
+            parenthesized!(content in input);
+            if content.peek(syn::LitStr) {
+                let fmt = content.parse()?;
+                let args = if content.is_empty() {
+                    TokenStream::new()
+                } else {
+                    fmt::parse_token_expr(&content, false)?
+                };
+                let display = Display {
+                    fmt,
+                    args,
+                    has_bonus_display: false,
+                };
+                Some(display)
+            } else {
+                return Err(syn::Error::new(input.span(), "Invalid argument to label() attribute. The first argument must be a literal string."));
+            }
+        } else if la.peek(Token![=]) {
+            // #[label = "blabla"]
+            input.parse::<Token![=]>()?;
+            Some(Display {
+                fmt: input.parse()?,
+                args: TokenStream::new(),
+                has_bonus_display: false,
+            })
+        } else {
+            None
+        };
+        Ok(LabelAttr { label })
+    }
+}
+
+impl Labels {
+    pub fn from_fields(fields: &syn::Fields) -> syn::Result<Option<Self>> {
+        match fields {
+            syn::Fields::Named(named) => Self::from_fields_vec(named.named.iter().collect()),
+            syn::Fields::Unnamed(unnamed) => {
+                Self::from_fields_vec(unnamed.unnamed.iter().collect())
+            }
+            syn::Fields::Unit => Ok(None),
+        }
+    }
+
+    fn from_fields_vec(fields: Vec<&syn::Field>) -> syn::Result<Option<Self>> {
+        let mut labels = Vec::new();
+        for (i, field) in fields.iter().enumerate() {
+            for attr in &field.attrs {
+                if attr.path.is_ident("label") {
+                    let span = if let Some(ident) = field.ident.clone() {
+                        syn::Member::Named(ident)
+                    } else {
+                        syn::Member::Unnamed(syn::Index {
+                            index: i as u32,
+                            span: field.span(),
+                        })
+                    };
+                    let LabelAttr { label } = syn::parse2::<LabelAttr>(attr.tokens.clone())?;
+                    labels.push(Label { label, span });
+                }
+            }
+        }
+        if labels.is_empty() {
+            Ok(None)
+        } else {
+            Ok(Some(Labels(labels)))
+        }
+    }
+
+    pub(crate) fn gen_struct(&self, fields: &syn::Fields) -> Option<TokenStream> {
+        let (display_pat, display_members) = display_pat_members(fields);
+        let labels = self.0.iter().map(|highlight| {
+            let Label { span, label } = highlight;
+            if let Some(display) = label {
+                let (fmt, args) = display.expand_shorthand_cloned(&display_members);
+                quote! {
+                    miette::LabeledSpan::new_with_span(
+                        std::option::Option::Some(format!(#fmt #args)),
+                        self.#span.clone(),
+                    )
+                }
+            } else {
+                quote! {
+                    miette::LabeledSpan::new_with_span(
+                        std::option::Option::None,
+                        self.#span.clone(),
+                    )
+                }
+            }
+        });
+        Some(quote! {
+            #[allow(unused_variables)]
+            fn labels(&self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = miette::LabeledSpan> + '_>> {
+                let Self #display_pat = self;
+                Some(Box::new(vec![
+                    #(#labels),*
+                ].into_iter()))
+            }
+        })
+    }
+
+    pub(crate) fn gen_enum(variants: &[DiagnosticDef]) -> Option<TokenStream> {
+        gen_all_variants_with(
+            variants,
+            WhichFn::Labels,
+            |ident, fields, DiagnosticConcreteArgs { labels, .. }| {
+                let (display_pat, display_members) = display_pat_members(fields);
+                labels.as_ref().and_then(|labels| {
+                let variant_labels = labels.0.iter().map(|label| {
+                    let Label { span, label } = label;
+                    let field = match &span {
+                        syn::Member::Named(ident) => ident.clone(),
+                        syn::Member::Unnamed(syn::Index { index, .. }) => {
+                            format_ident!("_{}", index)
+                        }
+                    };
+                    if let Some(display) = label {
+                        let (fmt, args) = display.expand_shorthand_cloned(&display_members);
+                        quote! {
+                            miette::LabeledSpan::new_with_span(
+                                std::option::Option::Some(format!(#fmt #args)),
+                                #field.clone(),
+                            )
+                        }
+                    } else {
+                        quote! {
+                            miette::LabeledSpan::new_with_span(
+                                std::option::Option::None,
+                                #field.clone(),
+                            )
+                        }
+                    }
+                });
+                let variant_name = ident.clone();
+                match &fields {
+                    syn::Fields::Unit => None,
+                    _ => Some(quote! {
+                        Self::#variant_name #display_pat => std::option::Option::Some(std::boxed::Box::new(vec![
+                            #(#variant_labels),*
+                        ].into_iter())),
+                    }),
+                }
+                })
+            },
+        )
+    }
+}
diff --git a/miette-derive/src/lib.rs b/miette-derive/src/lib.rs
index c99f813..c65b3f5 100644
--- a/miette-derive/src/lib.rs
+++ b/miette-derive/src/lib.rs
@@ -9,12 +9,13 @@ mod diagnostic_arg;
 mod fmt;
 mod forward;
 mod help;
+mod label;
 mod severity;
-mod snippets;
+mod source_code;
 mod url;
 mod utils;
 
-#[proc_macro_derive(Diagnostic, attributes(diagnostic, snippet, highlight))]
+#[proc_macro_derive(Diagnostic, attributes(diagnostic, label, source_code))]
 pub fn derive_diagnostic(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let input = parse_macro_input!(input as DeriveInput);
     let cmd = match Diagnostic::from_derive_input(input) {
diff --git a/miette-derive/src/snippets.rs b/miette-derive/src/snippets.rs
deleted file mode 100644
index 26a53b1..0000000
--- a/miette-derive/src/snippets.rs
+++ /dev/null
@@ -1,375 +0,0 @@
-use std::collections::HashMap;
-
-use proc_macro2::TokenStream;
-use quote::{format_ident, quote};
-use syn::{
-    parenthesized,
-    parse::{Parse, ParseStream},
-    spanned::Spanned,
-    Token,
-};
-
-use crate::{
-    diagnostic::{DiagnosticConcreteArgs, DiagnosticDef},
-    fmt::{self, Display},
-    forward::WhichFn,
-    utils::{display_pat_members, gen_all_variants_with},
-};
-
-pub struct Snippets(Vec<Snippet>);
-
-struct Snippet {
-    message: Option<Display>,
-    highlights: Vec<Highlight>,
-    source: syn::Member,
-    snippet: syn::Member,
-}
-
-struct Highlight {
-    label: Option<Display>,
-    highlight: syn::Member,
-}
-
-struct SnippetAttr {
-    source: syn::Member,
-    message: Option<Display>,
-}
-
-struct HighlightAttr {
-    label: Option<Display>,
-    snippet: syn::Member,
-}
-
-impl Parse for SnippetAttr {
-    fn parse(input: ParseStream) -> syn::Result<Self> {
-        let source = input.parse::<syn::Member>()?;
-        let message = if input.peek(Token![,]) {
-            input.parse::<Token![,]>()?;
-            let ident = input.parse::<syn::Ident>()?;
-            if ident == "message" {
-                let la = input.lookahead1();
-                if la.peek(syn::token::Paren) {
-                    let content;
-                    parenthesized!(content in input);
-                    if content.peek(syn::LitStr) {
-                        let fmt = content.parse()?;
-                        let args = if content.is_empty() {
-                            TokenStream::new()
-                        } else {
-                            fmt::parse_token_expr(&content, false)?
-                        };
-                        let display = Display {
-                            fmt,
-                            args,
-                            has_bonus_display: false,
-                        };
-                        Some(display)
-                    } else {
-                        return Err(syn::Error::new(ident.span(), "Invalid argument to message() sub-attribute. The first argument must be a literal string."));
-                    }
-                } else {
-                    input.parse::<Token![=]>()?;
-                    Some(Display {
-                        fmt: input.parse()?,
-                        args: TokenStream::new(),
-                        has_bonus_display: false,
-                    })
-                }
-            } else {
-                return Err(syn::Error::new(
-                    ident.span(),
-                    "Invalid sub-attribute. Only `message()` is allowed.",
-                ));
-            }
-        } else {
-            None
-        };
-        Ok(SnippetAttr { source, message })
-    }
-}
-
-impl Parse for HighlightAttr {
-    fn parse(input: ParseStream) -> syn::Result<Self> {
-        let snippet = input.parse::<syn::Member>()?;
-        let label = if input.peek(Token![,]) {
-            input.parse::<Token![,]>()?;
-            let ident = input.parse::<syn::Ident>()?;
-            if ident == "label" {
-                let la = input.lookahead1();
-                if la.peek(syn::token::Paren) {
-                    let content;
-                    parenthesized!(content in input);
-                    if content.peek(syn::LitStr) {
-                        let fmt = content.parse()?;
-                        let args = if content.is_empty() {
-                            TokenStream::new()
-                        } else {
-                            fmt::parse_token_expr(&content, false)?
-                        };
-                        let display = Display {
-                            fmt,
-                            args,
-                            has_bonus_display: false,
-                        };
-                        Some(display)
-                    } else {
-                        return Err(syn::Error::new(ident.span(), "Invalid argument to label() sub-attribute. The first argument must be a literal string."));
-                    }
-                } else {
-                    input.parse::<Token![=]>()?;
-                    Some(Display {
-                        fmt: input.parse()?,
-                        args: TokenStream::new(),
-                        has_bonus_display: false,
-                    })
-                }
-            } else {
-                return Err(syn::Error::new(
-                    ident.span(),
-                    "Invalid sub-attribute. Only `label()` is allowed.",
-                ));
-            }
-        } else {
-            None
-        };
-        Ok(HighlightAttr { snippet, label })
-    }
-}
-
-impl Snippets {
-    pub fn from_fields(fields: &syn::Fields) -> syn::Result<Option<Self>> {
-        match fields {
-            syn::Fields::Named(named) => Self::from_fields_vec(named.named.iter().collect()),
-            syn::Fields::Unnamed(unnamed) => {
-                Self::from_fields_vec(unnamed.unnamed.iter().collect())
-            }
-            syn::Fields::Unit => Ok(None),
-        }
-    }
-
-    fn from_fields_vec(fields: Vec<&syn::Field>) -> syn::Result<Option<Self>> {
-        let mut snippets = HashMap::new();
-        // First we collect all the contexts
-        for (i, field) in fields.iter().enumerate() {
-            for attr in &field.attrs {
-                if attr.path.is_ident("snippet") {
-                    let snippet = if let Some(ident) = field.ident.clone() {
-                        syn::Member::Named(ident)
-                    } else {
-                        syn::Member::Unnamed(syn::Index {
-                            index: i as u32,
-                            span: field.span(),
-                        })
-                    };
-                    let SnippetAttr { source, message } = attr.parse_args::<SnippetAttr>()?;
-                    // TODO: useful error when source refers to a field that doesn't exist.
-                    snippets.insert(
-                        snippet.clone(),
-                        Snippet {
-                            message,
-                            highlights: Vec::new(),
-                            source,
-                            snippet,
-                        },
-                    );
-                }
-            }
-        }
-        // Then we loop again looking for highlights
-        for (i, field) in fields.iter().enumerate() {
-            for attr in &field.attrs {
-                if attr.path.is_ident("highlight") {
-                    let HighlightAttr { snippet, label } = attr.parse_args::<HighlightAttr>()?;
-                    if let Some(snippet) = snippets.get_mut(&snippet) {
-                        let member = if let Some(ident) = field.ident.clone() {
-                            syn::Member::Named(ident)
-                        } else {
-                            syn::Member::Unnamed(syn::Index {
-                                index: i as u32,
-                                span: field.span(),
-                            })
-                        };
-                        snippet.highlights.push(Highlight {
-                            highlight: member,
-                            label,
-                        });
-                    } else {
-                        return Err(syn::Error::new(snippet.span(), "Highlight must refer to an existing field with a #[snippet(...)] attribute."));
-                    }
-                }
-            }
-        }
-        if snippets.is_empty() {
-            Ok(None)
-        } else {
-            Ok(Some(Snippets(snippets.into_values().collect())))
-        }
-    }
-
-    pub(crate) fn gen_struct(&self, fields: &syn::Fields) -> Option<TokenStream> {
-        let (display_pat, display_members) = display_pat_members(fields);
-        let snippets = self.0.iter().map(|snippet| {
-            // snippet message
-            let msg = if let Some(display) = &snippet.message {
-                let (fmt, args) = display.expand_shorthand_cloned(&display_members);
-                quote! {
-                    message: {
-                        std::option::Option::Some(format!(#fmt #args))
-                    },
-                }
-            } else {
-                quote! {
-                    message: std::option::Option::None,
-                }
-            };
-
-            // Source field
-            let src_ident = &snippet.source;
-            let src_ident = quote! {
-                source: &self.#src_ident,
-            };
-
-            // Context
-            let context = &snippet.snippet;
-            let context = quote! {
-                context: self.#context.clone().into(),
-            };
-
-            // Highlights
-            let highlights = snippet.highlights.iter().map(|highlight| {
-                let Highlight { highlight, label } = highlight;
-                if let Some(display) = label {
-                    let (fmt, args) = display.expand_shorthand_cloned(&display_members);
-                    quote! {
-                        (
-                            std::option::Option::Some(format!(#fmt #args)),
-                            self.#highlight.clone().into()
-                        )
-                    }
-                } else {
-                    quote! {
-                        (std::option::Option::None, self.#highlight.clone().into())
-                    }
-                }
-            });
-            let highlights = quote! {
-                highlights: std::option::Option::Some(vec![
-                    #(#highlights),*
-                ]),
-            };
-
-            // Generate the snippet itself
-            quote! {
-                miette::DiagnosticSnippet {
-                    #msg
-                    #src_ident
-                    #context
-                    #highlights
-                }
-            }
-        });
-        Some(quote! {
-            #[allow(unused_variables)]
-            fn snippets(&self) -> std::option::Option<std::boxed::Box<dyn std::iter::Iterator<Item = miette::DiagnosticSnippet> + '_>> {
-                let Self #display_pat = self;
-                Some(Box::new(vec![
-                    #(#snippets),*
-                ].into_iter()))
-            }
-        })
-    }
-
-    pub(crate) fn gen_enum(variants: &[DiagnosticDef]) -> Option<TokenStream> {
-        gen_all_variants_with(
-            variants,
-            WhichFn::Snippets,
-            |ident, fields, DiagnosticConcreteArgs { snippets, .. }| {
-                let (display_pat, display_members) = display_pat_members(fields);
-                snippets.as_ref().and_then(|snippets| {
-                        let variant_snippets = snippets.0.iter().map(|snippet| {
-                            // snippet message
-                            let msg = if let Some(display) = &snippet.message {
-                                let (fmt, args) = display.expand_shorthand_cloned(&display_members);
-                                quote! {
-                                    message: std::option::Option::Some(format!(#fmt, #args)),
-                                }
-                            } else {
-                                quote! {
-                                    message: std::option::Option::None,
-                                }
-                            };
-                            // Source field
-                            let src_ident = match &snippet.source {
-                                syn::Member::Named(id) => id.clone(),
-                                syn::Member::Unnamed(syn::Index { index, .. }) => {
-                                    format_ident!("_{}", index)
-                                }
-                            };
-                            let src_ident = quote! {
-                                // TODO: I don't like this. Think about it more and maybe improve protocol?
-                                source: #src_ident,
-                            };
-
-                            // Context
-                            let context = match &snippet.snippet {
-                                syn::Member::Named(id) => id.clone(),
-                                syn::Member::Unnamed(syn::Index { index, .. }) => {
-                                    format_ident!("_{}", index)
-                                }
-                            };
-                            let context = quote! {
-                                context: #context.clone().into(),
-                            };
-
-                            // Highlights
-                            let highlights = snippet.highlights.iter().map(|highlight| {
-                                let Highlight { highlight, label } = highlight;
-                                let m = match highlight {
-                                    syn::Member::Named(id) => id.clone(),
-                                    syn::Member::Unnamed(syn::Index { index, .. }) => {
-                                        format_ident!("_{}", index)
-                                    }
-                                };
-                                if let Some(Display { fmt, args, ..}) = label {
-                                    quote! {
-                                        (
-                                            std::option::Option::Some(format!(#fmt, #args)),
-                                            #m.clone().into()
-                                        )
-                                    }
-                                } else {
-                                    quote! {
-                                        (std::option::Option::None, #m.clone().into())
-                                    }
-                                }
-                            });
-                            let highlights = quote! {
-                                highlights: std::option::Option::Some(vec![
-                                    #(#highlights),*
-                                ]),
-                            };
-
-                            // Generate the snippet itself
-                            quote! {
-                                miette::DiagnosticSnippet {
-                                    #msg
-                                    #src_ident
-                                    #context
-                                    #highlights
-                                }
-                            }
-                        });
-                        let variant_name = ident.clone();
-                        match &fields {
-                            syn::Fields::Unit => None,
-                            _ => Some(quote! {
-                                Self::#variant_name #display_pat => std::option::Option::Some(std::boxed::Box::new(vec![
-                                    #(#variant_snippets),*
-                                ].into_iter())),
-                            }),
-                        }
-                    })
-            },
-        )
-    }
-}
diff --git a/miette-derive/src/source_code.rs b/miette-derive/src/source_code.rs
new file mode 100644
index 0000000..dd80869
--- /dev/null
+++ b/miette-derive/src/source_code.rs
@@ -0,0 +1,81 @@
+use proc_macro2::TokenStream;
+use quote::{format_ident, quote};
+use syn::spanned::Spanned;
+
+use crate::{
+    diagnostic::{DiagnosticConcreteArgs, DiagnosticDef},
+    forward::WhichFn,
+    utils::{display_pat_members, gen_all_variants_with},
+};
+
+pub struct SourceCode {
+    source_code: syn::Member,
+}
+
+impl SourceCode {
+    pub fn from_fields(fields: &syn::Fields) -> syn::Result<Option<Self>> {
+        match fields {
+            syn::Fields::Named(named) => Self::from_fields_vec(named.named.iter().collect()),
+            syn::Fields::Unnamed(unnamed) => {
+                Self::from_fields_vec(unnamed.unnamed.iter().collect())
+            }
+            syn::Fields::Unit => Ok(None),
+        }
+    }
+
+    fn from_fields_vec(fields: Vec<&syn::Field>) -> syn::Result<Option<Self>> {
+        for (i, field) in fields.iter().enumerate() {
+            for attr in &field.attrs {
+                if attr.path.is_ident("source_code") {
+                    let source_code = if let Some(ident) = field.ident.clone() {
+                        syn::Member::Named(ident)
+                    } else {
+                        syn::Member::Unnamed(syn::Index {
+                            index: i as u32,
+                            span: field.span(),
+                        })
+                    };
+                    return Ok(Some(SourceCode { source_code }));
+                }
+            }
+        }
+        Ok(None)
+    }
+
+    pub(crate) fn gen_struct(&self, fields: &syn::Fields) -> Option<TokenStream> {
+        let (display_pat, _display_members) = display_pat_members(fields);
+        let src = &self.source_code;
+        Some(quote! {
+            #[allow(unused_variables)]
+            fn source_code(&self) -> std::option::Option<&dyn miette::SourceCode> {
+                let Self #display_pat = self;
+                Some(&self.#src)
+            }
+        })
+    }
+
+    pub(crate) fn gen_enum(variants: &[DiagnosticDef]) -> Option<TokenStream> {
+        gen_all_variants_with(
+            variants,
+            WhichFn::SourceCode,
+            |ident, fields, DiagnosticConcreteArgs { source_code, .. }| {
+                let (display_pat, _display_members) = display_pat_members(fields);
+                source_code.as_ref().and_then(|source_code| {
+                    let field = match &source_code.source_code {
+                        syn::Member::Named(ident) => ident.clone(),
+                        syn::Member::Unnamed(syn::Index { index, .. }) => {
+                            format_ident!("_{}", index)
+                        }
+                    };
+                    let variant_name = ident.clone();
+                    match &fields {
+                        syn::Fields::Unit => None,
+                        _ => Some(quote! {
+                            Self::#variant_name #display_pat => std::option::Option::Some(#field),
+                        }),
+                    }
+                })
+            },
+        )
+    }
+}
diff --git a/src/eyreish/context.rs b/src/eyreish/context.rs
index 4d16fc3..e2547a6 100644
--- a/src/eyreish/context.rs
+++ b/src/eyreish/context.rs
@@ -4,7 +4,7 @@ use core::fmt::{self, Debug, Display, Write};
 
 use std::error::Error as StdError;
 
-use crate::Diagnostic;
+use crate::{Diagnostic, LabeledSpan};
 
 mod ext {
     use super::*;
@@ -141,10 +141,8 @@ where
         self.error.url()
     }
 
-    fn snippets<'a>(
-        &'a self,
-    ) -> Option<Box<dyn Iterator<Item = crate::DiagnosticSnippet<'a>> + 'a>> {
-        self.error.snippets()
+    fn labels<'a>(&'a self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + 'a>> {
+        self.error.labels()
     }
 }
 
@@ -168,10 +166,8 @@ where
         self.error.inner.diagnostic().url()
     }
 
-    fn snippets<'a>(
-        &'a self,
-    ) -> Option<Box<dyn Iterator<Item = crate::DiagnosticSnippet<'a>> + 'a>> {
-        self.error.inner.diagnostic().snippets()
+    fn labels<'a>(&'a self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + 'a>> {
+        self.error.inner.diagnostic().labels()
     }
 }
 
diff --git a/src/eyreish/wrapper.rs b/src/eyreish/wrapper.rs
index 3104dad..e58ce4d 100644
--- a/src/eyreish/wrapper.rs
+++ b/src/eyreish/wrapper.rs
@@ -2,7 +2,7 @@ use core::fmt::{self, Debug, Display};
 
 use std::error::Error as StdError;
 
-use crate::Diagnostic;
+use crate::{Diagnostic, LabeledSpan};
 
 use crate as miette;
 
@@ -71,11 +71,31 @@ impl Display for NoneError {
 impl StdError for NoneError {}
 impl Diagnostic for NoneError {}
 
-#[derive(miette_derive::Diagnostic)]
 #[repr(transparent)]
-#[diagnostic(transparent)]
 pub(crate) struct BoxedError(pub(crate) Box<dyn Diagnostic + Send + Sync>);
 
+impl Diagnostic for BoxedError {
+    fn code<'a>(&'a self) -> Option<Box<dyn Display + 'a>> {
+        self.0.code()
+    }
+
+    fn severity(&self) -> Option<miette::Severity> {
+        self.0.severity()
+    }
+
+    fn help<'a>(&'a self) -> Option<Box<dyn Display + 'a>> {
+        self.0.help()
+    }
+
+    fn url<'a>(&'a self) -> Option<Box<dyn Display + 'a>> {
+        self.0.url()
+    }
+
+    fn labels<'a>(&'a self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + 'a>> {
+        self.0.labels()
+    }
+}
+
 impl Debug for BoxedError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         Debug::fmt(&self.0, f)
diff --git a/src/handlers/debug.rs b/src/handlers/debug.rs
index 2248832..e6f0733 100644
--- a/src/handlers/debug.rs
+++ b/src/handlers/debug.rs
@@ -48,9 +48,9 @@ impl DebugReportHandler {
         if let Some(help) = diagnostic.help() {
             diag.field("help", &help.to_string());
         }
-        if let Some(snippets) = diagnostic.snippets() {
-            let snippets: Vec<_> = snippets.collect();
-            diag.field("snippets", &format!("{:?}", snippets));
+        if let Some(labels) = diagnostic.labels() {
+            let labels: Vec<_> = labels.collect();
+            diag.field("labels", &format!("{:?}", labels));
         }
         diag.finish()?;
         writeln!(f)?;
diff --git a/src/handlers/graphical.rs b/src/handlers/graphical.rs
index 7e52b9c..a960fc4 100644
--- a/src/handlers/graphical.rs
+++ b/src/handlers/graphical.rs
@@ -1,12 +1,13 @@
 use std::fmt::{self, Write};
 
+use itertools::Itertools;
 use owo_colors::{OwoColorize, Style};
 use unicode_width::UnicodeWidthStr;
 
 use crate::chain::Chain;
 use crate::handlers::theme::*;
-use crate::protocol::{Diagnostic, DiagnosticSnippet, Severity};
-use crate::{ReportHandler, SourceSpan, SpanContents};
+use crate::protocol::{Diagnostic, Severity};
+use crate::{LabeledSpan, ReportHandler, SourceCode, SourceSpan, SpanContents};
 
 /**
 A [ReportHandler] that displays a given [crate::Report] in a quasi-graphical
@@ -94,10 +95,14 @@ impl GraphicalReportHandler {
         writeln!(f)?;
         self.render_causes(f, diagnostic)?;
 
-        if let Some(snippets) = diagnostic.snippets() {
-            for snippet in snippets {
-                writeln!(f)?;
-                self.render_snippet(f, &snippet)?;
+        if let Some(source) = diagnostic.source_code() {
+            if let Some(labels) = diagnostic.labels() {
+                let mut labels = labels.collect::<Vec<_>>();
+                labels.sort_unstable_by_key(|l| l.inner().offset());
+                if !labels.is_empty() {
+                    writeln!(f)?;
+                    self.render_snippets(f, source, labels)?;
+                }
             }
         }
 
@@ -228,23 +233,39 @@ impl GraphicalReportHandler {
         Ok(())
     }
 
-    fn render_snippet(
+    fn render_snippets(
         &self,
         f: &mut impl fmt::Write,
-        snippet: &DiagnosticSnippet<'_>,
+        source: &dyn SourceCode,
+        labels: Vec<LabeledSpan>,
     ) -> fmt::Result {
-        let (contents, lines) = self.get_lines(snippet)?;
+        // TODO: Actually do the rewrite against the new protocol.
+        let contexts: Vec<_> = labels
+            .iter()
+            .cloned()
+            .coalesce(|left, right| {
+                if left.offset() + left.len() >= right.offset() {
+                    let left_end = left.offset() + left.len();
+                    let right_end = right.offset() + right.len();
+                    Ok(LabeledSpan::new(
+                        left.label().map(String::from),
+                        left.offset(),
+                        right_end - left_end,
+                    ))
+                } else {
+                    Err((left, right))
+                }
+            })
+            .collect();
+        let (contents, lines) = self.get_lines(source, &labels)?;
 
-        // Highlights are the bits we're going to underline in our overall
-        // snippet, and we need to do some analysis first to come up with
-        // gutter size.
-        let mut highlights = snippet.highlights.clone().unwrap_or_else(Vec::new);
-        // sorting is your friend.
-        highlights.sort_unstable_by_key(|(_, h)| h.offset());
-        let highlights = highlights
-            .into_iter()
+        // sorting is your friend
+        let labels = labels
+            .iter()
             .zip(self.theme.styles.highlights.iter().cloned().cycle())
-            .map(|((label, hl), st)| FancySpan::new(label, hl, st))
+            .map(|(label, st)| {
+                FancySpan::new(label.label().map(String::from), label.inner().clone(), st)
+            })
             .collect::<Vec<_>>();
 
         // The max number of gutter-lines that will be active at any given
@@ -253,7 +274,7 @@ impl GraphicalReportHandler {
         let mut max_gutter = 0usize;
         for line in &lines {
             let mut num_highlights = 0;
-            for hl in &highlights {
+            for hl in &labels {
                 if !line.span_line_only(hl) && line.span_applies(hl) {
                     num_highlights += 1;
                 }
@@ -270,45 +291,26 @@ impl GraphicalReportHandler {
             .len();
 
         // Header
-        if let Some(msg) = &snippet.message {
-            writeln!(
-                f,
-                "{}{}{}",
-                " ".repeat(linum_width + 2),
-                self.theme.characters.ltop,
-                self.theme.characters.hbar.to_string().repeat(4)
-            )?;
-            writeln!(
-                f,
-                "{}{} error: {}",
-                " ".repeat(linum_width + 2),
-                self.theme.characters.vbar,
-                msg
-            )?;
-        }
         write!(
             f,
             "{}{}{}",
             " ".repeat(linum_width + 2),
-            if snippet.message.is_some() {
-                self.theme.characters.lcross
-            } else {
-                self.theme.characters.ltop
-            },
+            self.theme.characters.ltop,
             self.theme.characters.hbar,
         )?;
-        if let Some(source_name) = snippet.source.name() {
-            let source_name = source_name.style(self.theme.styles.link);
-            writeln!(
-                f,
-                "[{}:{}:{}]",
-                source_name,
-                contents.line() + 1,
-                contents.column() + 1
-            )?;
-        } else {
-            writeln!(f, "[{}:{}]", contents.line() + 1, contents.column() + 1)?;
-        }
+        // TODO: filenames
+        // if let Some(source_name) = source.name() {
+        //     let source_name = source_name.style(self.theme.styles.link);
+        //     writeln!(
+        //         f,
+        //         "[{}:{}:{}]",
+        //         source_name,
+        //         contents.line() + 1,
+        //         contents.column() + 1
+        //     )?;
+        // } else {
+        //     writeln!(f, "[{}:{}]", contents.line() + 1, contents.column() + 1)?;
+        // }
 
         // Blank line to improve readability
         writeln!(
@@ -326,13 +328,13 @@ impl GraphicalReportHandler {
             // Then, we need to print the gutter, along with any fly-bys We
             // have separate gutters depending on whether we're on the actual
             // line, or on one of the "highlight lines" below it.
-            self.render_line_gutter(f, max_gutter, line, &highlights)?;
+            self.render_line_gutter(f, max_gutter, line, &labels)?;
 
             // And _now_ we can print out the line text itself!
             writeln!(f, "{}", line.text)?;
 
             // Next, we write all the highlights that apply to this particular line.
-            let (single_line, multi_line): (Vec<_>, Vec<_>) = highlights
+            let (single_line, multi_line): (Vec<_>, Vec<_>) = labels
                 .iter()
                 .filter(|hl| line.span_applies(hl))
                 .partition(|hl| line.span_line_only(hl));
@@ -340,14 +342,14 @@ impl GraphicalReportHandler {
                 // no line number!
                 self.write_no_linum(f, linum_width)?;
                 // gutter _again_
-                self.render_highlight_gutter(f, max_gutter, line, &highlights)?;
+                self.render_highlight_gutter(f, max_gutter, line, &labels)?;
                 self.render_single_line_highlights(
                     f,
                     line,
                     linum_width,
                     max_gutter,
                     &single_line,
-                    &highlights,
+                    &labels,
                 )?;
             }
             for hl in multi_line {
@@ -355,7 +357,7 @@ impl GraphicalReportHandler {
                     // no line number!
                     self.write_no_linum(f, linum_width)?;
                     // gutter _again_
-                    self.render_highlight_gutter(f, max_gutter, line, &highlights)?;
+                    self.render_highlight_gutter(f, max_gutter, line, &labels)?;
                     self.render_multi_line_end(f, hl)?;
                 }
             }
@@ -571,16 +573,23 @@ impl GraphicalReportHandler {
 
     fn get_lines<'a>(
         &'a self,
-        snippet: &'a DiagnosticSnippet<'_>,
-    ) -> Result<(Box<dyn SpanContents + 'a>, Vec<Line>), fmt::Error> {
-        let context_data = snippet
-            .source
-            .read_span(&snippet.context)
+        source: &'a dyn SourceCode,
+        labels: &'a [LabeledSpan],
+    ) -> Result<(Box<dyn SpanContents<'a> + 'a>, Vec<Line>), fmt::Error> {
+        let first = labels.first().expect("MIETTE BUG: This should be safe.");
+        let last = labels.last().expect("MIETTE BUG: This should be safe.");
+        let context_span = (
+            first.inner().offset(),
+            last.inner().offset() + last.inner().len(),
+        )
+            .into();
+        let context_data = source
+            .read_span(&context_span, 1, 1)
             .map_err(|_| fmt::Error)?;
         let context = std::str::from_utf8(context_data.data()).expect("Bad utf8 detected");
         let mut line = context_data.line();
         let mut column = context_data.column();
-        let mut offset = snippet.context.offset();
+        let mut offset = context_span.offset();
         let mut line_offset = offset;
         let mut iter = context.chars().peekable();
         let mut line_str = String::new();
diff --git a/src/handlers/narratable.rs b/src/handlers/narratable.rs
index 8482330..31ec81c 100644
--- a/src/handlers/narratable.rs
+++ b/src/handlers/narratable.rs
@@ -1,8 +1,8 @@
 use std::fmt;
 
 use crate::chain::Chain;
-use crate::protocol::{Diagnostic, DiagnosticSnippet, Severity};
-use crate::{ReportHandler, SourceSpan, SpanContents};
+use crate::protocol::{Diagnostic, Severity};
+use crate::{ReportHandler, SourceCode, SourceSpan, SpanContents};
 
 /**
 [ReportHandler] that renders plain text and avoids extraneous graphics.
@@ -18,7 +18,7 @@ impl NarratableReportHandler {
     /// Create a new [NarratableReportHandler]. There are no customization
     /// options.
     pub fn new() -> Self {
-        Self { footer: None}
+        Self { footer: None }
     }
 
     /// Set the footer to be displayed at the end of the report.
@@ -47,12 +47,11 @@ impl NarratableReportHandler {
         self.render_header(f, diagnostic)?;
         self.render_causes(f, diagnostic)?;
 
-        if let Some(snippets) = diagnostic.snippets() {
-            for snippet in snippets {
-                writeln!(f)?;
-                self.render_snippet(f, &snippet)?;
-            }
-        }
+        // if let Some(labels) = diagnostic.labels() {
+        //     for label in labels {
+        //         self.render_label(f, &label)?;
+        //     }
+        // }
 
         self.render_footer(f, diagnostic)?;
         Ok(())
@@ -92,68 +91,12 @@ impl NarratableReportHandler {
         Ok(())
     }
 
-    fn render_snippet(
-        &self,
-        f: &mut impl fmt::Write,
-        snippet: &DiagnosticSnippet<'_>,
-    ) -> fmt::Result {
-        let (contents, lines) = self.get_lines(snippet)?;
-
-        write!(f, "Begin snippet")?;
-        if let Some(filename) = snippet.source.name() {
-            write!(f, " for {}", filename,)?;
-        }
-        write!(
-            f,
-            " starting at line {}, column {}",
-            contents.line() + 1,
-            contents.column() + 1
-        )?;
-        if let Some(message) = snippet.message.as_deref() {
-            write!(f, ": {}", message)?;
-        }
-        writeln!(f)?;
-        writeln!(f)?;
-
-        // Highlights are the bits we're going to underline in our overall
-        // snippet, and we need to do some analysis first to come up with
-        // gutter size.
-        let mut highlights = snippet.highlights.clone().unwrap_or_else(Vec::new);
-        // sorting is your friend.
-        highlights.sort_unstable_by_key(|(_, h)| h.offset());
-
-        // Now it's time for the fun part--actually rendering everything!
-        for line in &lines {
-            writeln!(f, "snippet line {}: {}", line.line_number, line.text)?;
-            let relevant = highlights.iter().filter(|(_, hl)| line.span_starts(hl));
-            for (label, hl) in relevant {
-                let contents = snippet.source.read_span(hl).map_err(|_| fmt::Error)?;
-                if contents.line() + 1 == line.line_number {
-                    write!(
-                        f,
-                        "    highlight starting at line {}, column {}",
-                        contents.line() + 1,
-                        contents.column() + 1
-                    )?;
-                    if let Some(label) = label {
-                        write!(f, ": {}", label)?;
-                    }
-                    writeln!(f)?;
-                }
-            }
-        }
-        writeln!(f)?;
-        Ok(())
-    }
-
+    /*
     fn get_lines<'a>(
         &'a self,
-        snippet: &'a DiagnosticSnippet<'a>,
+        source: &'a dyn Source,
     ) -> Result<(Box<dyn SpanContents + 'a>, Vec<Line>), fmt::Error> {
-        let context_data = snippet
-            .source
-            .read_span(&snippet.context)
-            .map_err(|_| fmt::Error)?;
+        let context_data = source.read_span(&snippet.context).map_err(|_| fmt::Error)?;
         let context = std::str::from_utf8(context_data.data()).expect("Bad utf8 detected");
         let mut line = context_data.line();
         let mut column = context_data.column();
@@ -200,6 +143,7 @@ impl NarratableReportHandler {
         }
         Ok((context_data, lines))
     }
+    */
 }
 
 impl ReportHandler for NarratableReportHandler {
diff --git a/src/named_source.rs b/src/named_source.rs
index 0a6270d..656cc3d 100644
--- a/src/named_source.rs
+++ b/src/named_source.rs
@@ -1,38 +1,52 @@
-use crate::Source;
+use crate::{MietteError, MietteSpanContents, SourceCode, SpanContents};
 
 /// Utility struct for when you have a regular [Source] type, such as a String,
 /// that doesn't implement `name`, or if you want to override the `.name()`
 /// returned by the `Source`.
-#[derive(Debug)]
 pub struct NamedSource {
-    source: Box<dyn Source + Send + Sync + 'static>,
+    source: Box<dyn SourceCode + 'static>,
     name: String,
 }
 
+impl std::fmt::Debug for NamedSource {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("NamedSource")
+            .field("name", &self.name)
+            .field("source", &"<redacted>");
+        Ok(())
+    }
+}
+
 impl NamedSource {
-    /// Create a new [NamedSource] using a regular [Source] and giving it a [Source::name].
-    pub fn new(name: impl AsRef<str>, source: impl Source + Send + Sync + 'static) -> Self {
+    /// Create a new [NamedSource] using a regular [SourceCode] and giving its returned [SpanContents] a name.
+    pub fn new(name: impl AsRef<str>, source: impl SourceCode + Send + Sync + 'static) -> Self {
         Self {
             source: Box::new(source),
             name: name.as_ref().to_string(),
         }
     }
 
-    /// Returns a reference the inner [Source] type for this [NamedSource].
-    pub fn inner(&self) -> &(dyn Source + Send + Sync + 'static) {
+    /// Returns a reference the inner [SourceCode] type for this [NamedSource].
+    pub fn inner(&self) -> &(dyn SourceCode + 'static) {
         &*self.source
     }
 }
 
-impl Source for NamedSource {
+impl SourceCode for NamedSource {
     fn read_span<'a>(
         &'a self,
         span: &crate::SourceSpan,
-    ) -> Result<Box<dyn crate::SpanContents + 'a>, crate::MietteError> {
-        self.source.read_span(span)
-    }
-
-    fn name(&self) -> Option<String> {
-        Some(self.name.clone())
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        let contents = self
+            .inner()
+            .read_span(span, context_lines_before, context_lines_after)?;
+        Ok(Box::new(MietteSpanContents::new_named(
+            self.name.clone(),
+            contents.data(),
+            contents.line(),
+            contents.column(),
+        )))
     }
 }
diff --git a/src/protocol.rs b/src/protocol.rs
index f642def..3722367 100644
--- a/src/protocol.rs
+++ b/src/protocol.rs
@@ -45,9 +45,13 @@ pub trait Diagnostic: std::error::Error {
         None
     }
 
-    /// Additional contextual snippets. This is typically used for adding
-    /// marked-up source file output the way compilers often do.
-    fn snippets<'a>(&'a self) -> Option<Box<dyn Iterator<Item = DiagnosticSnippet<'a>> + 'a>> {
+    /// Source code to apply this Diagnostic's [Diagnostic::labels] to.
+    fn source_code(&self) -> Option<&dyn SourceCode> {
+        None
+    }
+
+    /// Labels to apply to this Diagnostic's [Diagnostic::source_code]
+    fn labels(&self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + '_>> {
         None
     }
 }
@@ -158,7 +162,7 @@ pub enum Severity {
 }
 
 /**
-Represents a readable source of some sort.
+Represents readable source code of some sort.
 
 This trait is able to support simple Source types like [String]s, as well
 as more involved types like indexes into centralized `SourceMap`-like types,
@@ -166,32 +170,86 @@ file handles, and even network streams.
 
 If you can read it, you can source it,
 and it's not necessary to read the whole thing--meaning you should be able to
-support Sources which are gigabytes or larger in size.
+support SourceCodes which are gigabytes or larger in size.
 */
-pub trait Source: std::fmt::Debug + Send + Sync {
-    /// Read the bytes for a specific span from this Source.
+pub trait SourceCode {
+    /// Read the bytes for a specific span from this SourceCode, keeping a
+    /// certain number of lines before and after the span as context.
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError>;
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError>;
+}
 
-    /// Optional name, usually a filename, for this source.
-    fn name(&self) -> Option<String> {
-        None
+/**
+A labeled [SourceSpan].
+*/
+#[derive(Debug, Clone)]
+pub struct LabeledSpan {
+    label: Option<String>,
+    span: SourceSpan,
+}
+
+impl LabeledSpan {
+    /// Makes a new labeled span.
+    pub fn new(label: Option<String>, offset: ByteOffset, len: ByteOffset) -> Self {
+        Self {
+            label,
+            span: (offset, len).into(),
+        }
+    }
+
+    /// Makes a new labeled span using an existing span.
+    pub fn new_with_span(label: Option<String>, span: impl Into<SourceSpan>) -> Self {
+        Self {
+            label,
+            span: span.into(),
+        }
+    }
+
+    /// Gets the (optional) label string for this LabeledSpan.
+    pub fn label(&self) -> Option<&str> {
+        self.label.as_deref()
+    }
+
+    /// Returns a reference to the inner [SourceSpan].
+    pub fn inner(&self) -> &SourceSpan {
+        &self.span
+    }
+
+    /// Returns the 0-based starting byte offset.
+    pub fn offset(&self) -> usize {
+        self.span.offset()
+    }
+
+    /// Returns the number of bytes this LabeledSpan spans.
+    pub fn len(&self) -> usize {
+        self.span.len()
+    }
+
+    /// True if this LabeledSpan is empty.
+    pub fn is_empty(&self) -> bool {
+        self.span.is_empty()
     }
 }
 
 /**
-Contents of a [Source] covered by [SourceSpan].
+Contents of a [SourceCode] covered by [SourceSpan].
 
 Includes line and column information to optimize highlight calculations.
 */
-pub trait SpanContents {
+pub trait SpanContents<'a> {
     /// Reference to the data inside the associated span, in bytes.
-    fn data(&self) -> &[u8];
-    /// The 0-indexed line in the associated [Source] where the data begins.
+    fn data(&self) -> &'a [u8];
+    /// An optional (file?) name for the container of this SpanContents.
+    fn name(&self) -> Option<&'a str> {
+        None
+    }
+    /// The 0-indexed line in the associated [SourceCode] where the data begins.
     fn line(&self) -> usize;
-    /// The 0-indexed column in the associated [Source] where the data begins,
+    /// The 0-indexed column in the associated [SourceCode] where the data begins,
     /// relative to `line`.
     fn column(&self) -> usize;
 }
@@ -201,23 +259,45 @@ Basic implementation of the [SpanContents] trait, for convenience.
 */
 #[derive(Clone, Debug)]
 pub struct MietteSpanContents<'a> {
-    /// Data from a [Source], in bytes.
+    /// Data from a [SourceCode], in bytes.
     data: &'a [u8],
     // The 0-indexed line where the associated [SourceSpan] _starts_.
     line: usize,
     // The 0-indexed column where the associated [SourceSpan] _starts_.
     column: usize,
+    // Optional filename
+    name: Option<String>,
 }
 
 impl<'a> MietteSpanContents<'a> {
     /// Make a new [MietteSpanContents] object.
     pub fn new(data: &'a [u8], line: usize, column: usize) -> MietteSpanContents<'a> {
-        MietteSpanContents { data, line, column }
+        MietteSpanContents {
+            data,
+            line,
+            column,
+            name: None,
+        }
+    }
+
+    /// Make a new [MietteSpanContents] object, with a name for its "file".
+    pub fn new_named(
+        name: String,
+        data: &'a [u8],
+        line: usize,
+        column: usize,
+    ) -> MietteSpanContents<'a> {
+        MietteSpanContents {
+            data,
+            line,
+            column,
+            name: Some(name),
+        }
     }
 }
 
-impl<'a> SpanContents for MietteSpanContents<'a> {
-    fn data(&self) -> &[u8] {
+impl<'a> SpanContents<'a> for MietteSpanContents<'a> {
+    fn data(&self) -> &'a [u8] {
         self.data
     }
     fn line(&self) -> usize {
@@ -229,24 +309,7 @@ impl<'a> SpanContents for MietteSpanContents<'a> {
 }
 
 /**
-A snippet from a [Source] to be displayed with a message and possibly some highlights.
- */
-#[derive(Clone, Debug)]
-pub struct DiagnosticSnippet<'a> {
-    /// Explanation of this specific diagnostic snippet.
-    pub message: Option<String>,
-    /// A [Source] that can be used to read the actual text of a source.
-    pub source: &'a (dyn Source),
-    /// The primary [SourceSpan] where this diagnostic is located.
-    pub context: SourceSpan,
-    /// Additional [SourceSpan]s that mark specific sections of the span, for
-    /// example, to underline specific text within the larger span. They're
-    /// paired with labels that should be applied to those sections.
-    pub highlights: Option<Vec<(Option<String>, SourceSpan)>>,
-}
-
-/**
-Span within a [Source] with an associated message.
+Span within a [SourceCode] with an associated message.
 */
 #[derive(Clone, Debug)]
 pub struct SourceSpan {
@@ -265,7 +328,7 @@ impl SourceSpan {
         }
     }
 
-    /// The absolute offset, in bytes, from the beginning of a [Source].
+    /// The absolute offset, in bytes, from the beginning of a [SourceCode].
     pub fn offset(&self) -> usize {
         self.offset.offset()
     }
@@ -301,12 +364,12 @@ impl From<(SourceOffset, SourceOffset)> for SourceSpan {
 }
 
 /**
-"Raw" type for the byte offset from the beginning of a [Source].
+"Raw" type for the byte offset from the beginning of a [SourceCode].
 */
 pub type ByteOffset = usize;
 
 /**
-Newtype that represents the [ByteOffset] from the beginning of a [Source]
+Newtype that represents the [ByteOffset] from the beginning of a [SourceCode]
 */
 #[derive(Clone, Copy, Debug)]
 pub struct SourceOffset(ByteOffset);
diff --git a/src/source_impls.rs b/src/source_impls.rs
index 8e5c8f7..fccab2b 100644
--- a/src/source_impls.rs
+++ b/src/source_impls.rs
@@ -1,5 +1,5 @@
 /*!
-Default trait implementations for [Source].
+Default trait implementations for [SourceCode].
 */
 use std::{
     borrow::{Cow, ToOwned},
@@ -7,98 +7,175 @@ use std::{
     sync::Arc,
 };
 
-use crate::{MietteError, MietteSpanContents, Source, SourceSpan, SpanContents};
+use crate::{MietteError, MietteSpanContents, SourceCode, SourceSpan, SpanContents};
 
-fn start_line_column(string: &str, span: &SourceSpan) -> Result<(usize, usize), MietteError> {
+fn context_info<'a>(
+    input: &'a [u8],
+    span: &SourceSpan,
+    context_lines_before: usize,
+    context_lines_after: usize,
+) -> Result<(&'a [u8], usize, usize), MietteError> {
     let mut offset = 0usize;
     let mut start_line = 0usize;
     let mut start_column = 0usize;
-    let mut iter = string.chars().peekable();
+    let mut before_lines_starts = Vec::new();
+    let mut current_line_start = 0usize;
+    let mut end_lines = 0usize;
+    let mut post_span = false;
+    let mut iter = input.iter().copied().peekable();
     while let Some(char) = iter.next() {
-        if offset < span.offset() {
-            match char {
-                '\r' => {
-                    if iter.next_if_eq(&'\n').is_some() {
+        if matches!(char, b'\r' | b'\n') {
+            if char == b'\r' && iter.next_if_eq(&b'\n').is_some() {
+                offset += 1;
+            }
+            if offset < span.offset() {
+                // We're before the start of the span.
+                start_column = 0;
+                before_lines_starts.push(current_line_start);
+                if before_lines_starts.len() > context_lines_before {
+                    start_line += 1;
+                    before_lines_starts.remove(0);
+                }
+            } else if offset >= span.offset() + span.len() - 1 {
+                // We're after the end of the span, but haven't necessarily
+                // started collecting end lines yet (we might still be
+                // collecting context lines).
+                if post_span {
+                    end_lines += 1;
+                    start_column = 0;
+                    if end_lines > context_lines_after {
                         offset += 1;
+                        break;
                     }
-                    start_line += 1;
-                    start_column = 0;
-                }
-                '\n' => {
-                    start_line += 1;
-                    start_column = 0;
-                }
-                _ => {
-                    start_column += 1;
                 }
             }
+            current_line_start = offset + 1;
+        } else if offset < span.offset() {
+            start_column += 1;
         }
 
         if offset >= span.offset() + span.len() - 1 {
-            return Ok((start_line, start_column));
+            post_span = true;
+            if end_lines >= context_lines_after {
+                offset += 1;
+                break;
+            }
         }
 
-        offset += char.len_utf8();
+        offset += 1;
+    }
+
+    if offset >= span.offset() + span.len() - 1 {
+        Ok((
+            &input[before_lines_starts
+                .get(0)
+                .copied()
+                .unwrap_or_else(|| span.offset())..offset],
+            start_line,
+            if context_lines_before == 0 {
+                start_column
+            } else {
+                0
+            },
+        ))
+    } else {
+        Err(MietteError::OutOfBounds)
     }
-    Err(MietteError::OutOfBounds)
 }
 
-// The basic impl here is on str (not &str), because otherwise String's impl cannot reuse it
-// without creating a temporary &str inside its read_span implementation, and then returning data
-// that refers to that temporary.
-impl Source for str {
+impl SourceCode for [u8] {
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError> {
-        let (start_line, start_column) = start_line_column(self, span)?;
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        let (data, start_line, start_column) =
+            context_info(self, span, context_lines_before, context_lines_after)?;
         return Ok(Box::new(MietteSpanContents::new(
-            &self.as_bytes()[span.offset()..span.offset() + span.len()],
+            data,
             start_line,
             start_column,
         )));
     }
 }
 
+impl<'src> SourceCode for &'src [u8] {
+    fn read_span<'a>(
+        &'a self,
+        span: &SourceSpan,
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        <[u8] as SourceCode>::read_span(self, span, context_lines_before, context_lines_after)
+    }
+}
+
+impl SourceCode for str {
+    fn read_span<'a>(
+        &'a self,
+        span: &SourceSpan,
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        <[u8] as SourceCode>::read_span(
+            self.as_bytes(),
+            span,
+            context_lines_before,
+            context_lines_after,
+        )
+    }
+}
+
 /// Makes `src: &'static str` or `struct S<'a> { src: &'a str }` usable.
-impl<'s> Source for &'s str {
+impl<'s> SourceCode for &'s str {
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError> {
-        <str as Source>::read_span(self, span)
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        <str as SourceCode>::read_span(self, span, context_lines_before, context_lines_after)
     }
 }
 
-impl Source for String {
+impl SourceCode for String {
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError> {
-        <str as Source>::read_span(self, span)
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        <str as SourceCode>::read_span(self, span, context_lines_before, context_lines_after)
     }
 }
 
-impl<T: Source> Source for Arc<T> {
+impl<T: SourceCode> SourceCode for Arc<T> {
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError> {
-        self.as_ref().read_span(span)
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        self.as_ref()
+            .read_span(span, context_lines_before, context_lines_after)
     }
 }
 
-impl<T: ?Sized + Source + ToOwned> Source for Cow<'_, T>
+impl<T: ?Sized + SourceCode + ToOwned> SourceCode for Cow<'_, T>
 where
-    // The minimal bounds are used here. `T::Owned` need not be `Source`, because `&T` can always
-    // be obtained from `Cow<'_, T>`.
+    // The minimal bounds are used here. `T::Owned` need not be `SourceCode`,
+    // because `&T` can always be obtained from `Cow<'_, T>`.
     T::Owned: Debug + Send + Sync,
 {
     fn read_span<'a>(
         &'a self,
         span: &SourceSpan,
-    ) -> Result<Box<dyn SpanContents + 'a>, MietteError> {
-        self.as_ref().read_span(span)
+        context_lines_before: usize,
+        context_lines_after: usize,
+    ) -> Result<Box<dyn SpanContents<'a> + 'a>, MietteError> {
+        self.as_ref()
+            .read_span(span, context_lines_before, context_lines_after)
     }
 }
 
@@ -109,24 +186,53 @@ mod tests {
     #[test]
     fn basic() -> Result<(), MietteError> {
         let src = String::from("foo\n");
-        let contents = src.read_span(&(0, 4).into())?;
+        let contents = src.read_span(&(0, 4).into(), 0, 0)?;
         assert_eq!("foo\n", std::str::from_utf8(contents.data()).unwrap());
+        assert_eq!(0, contents.line());
+        assert_eq!(0, contents.column());
         Ok(())
     }
 
     #[test]
     fn middle() -> Result<(), MietteError> {
         let src = String::from("foo\nbar\nbaz\n");
-        let contents = src.read_span(&(4, 4).into())?;
+        let contents = src.read_span(&(4, 4).into(), 0, 0)?;
         assert_eq!("bar\n", std::str::from_utf8(contents.data()).unwrap());
+        assert_eq!(1, contents.line());
+        assert_eq!(0, contents.column());
+        Ok(())
+    }
+
+    #[test]
+    fn middle_of_line() -> Result<(), MietteError> {
+        let src = String::from("foo\nbarbar\nbaz\n");
+        let contents = src.read_span(&(7, 4).into(), 0, 0)?;
+        assert_eq!("bar\n", std::str::from_utf8(contents.data()).unwrap());
+        assert_eq!(1, contents.line());
+        assert_eq!(3, contents.column());
         Ok(())
     }
 
     #[test]
     fn with_crlf() -> Result<(), MietteError> {
         let src = String::from("foo\r\nbar\r\nbaz\r\n");
-        let contents = src.read_span(&(5, 5).into())?;
+        let contents = src.read_span(&(5, 5).into(), 0, 0)?;
         assert_eq!("bar\r\n", std::str::from_utf8(contents.data()).unwrap());
+        assert_eq!(1, contents.line());
+        assert_eq!(0, contents.column());
+        Ok(())
+    }
+
+    #[test]
+    fn with_context() -> Result<(), MietteError> {
+        let src = String::from("xxx\nfoo\nbar\nbaz\n\nyyy\n");
+        let contents = src.read_span(&(8, 4).into(), 1, 2)?;
+        assert_eq!(
+            "foo\nbar\nbaz\n\n",
+            std::str::from_utf8(contents.data()).unwrap()
+        );
+        assert_eq!(1, contents.line());
+        assert_eq!(0, contents.column());
         Ok(())
     }
 }
diff --git a/tests/derive.rs b/tests/derive.rs
index 44a8827..2a7f354 100644
--- a/tests/derive.rs
+++ b/tests/derive.rs
@@ -202,41 +202,15 @@ fn test_snippet_named_struct() {
     #[error("welp")]
     #[diagnostic(code(foo::bar::baz))]
     struct Foo {
-        // The actual "source code" our contexts will be using. This can be
-        // reused by multiple contexts, and just needs to implement
-        // miette::Source!
+        #[source_code]
         src: String,
-
-        // The "snippet" span. This is the span that will be displayed to
-        // users. It should be a big enough slice of the Source to provide
-        // reasonable context, but still somewhat compact.
-        //
-        // You can have as many of these #[snippet] fields as you want, and
-        // even feed them from different sources!
-        //
-        // Example display:
-        //   / [my_snippet]: hi this is where the thing went wrong.
-        // 1 | hello
-        // 2 |     world
-        #[snippet(src, message("hi this is where the thing went wrong"))]
-        snip: SourceSpan, // Defines filename using `label`
-
-        // "Highlights" are the specific highlights _inside_ the snippet.
-        // These will be used to underline/point to specific sections of the
-        // #[snippet] they refer to. As such, these SourceSpans must be within
-        // the bounds of their referenced snippet.
-        //
-        // Example display:
-        // 1 | var1 + var2
-        //   | ^^^^   ^^^^ - var 2
-        //   | |
-        //   | var 1
-        #[highlight(snip)]
-        // label from SourceSpan is used, if any.
+        #[label("var 1")]
         var1: SourceSpan,
-        #[highlight(snip)]
+        #[label = "var 2"]
         // Anything that's Clone + Into<SourceSpan> can be used here.
         var2: (usize, usize),
+        #[label]
+        var3: (usize, usize),
     }
 }
 
@@ -246,15 +220,10 @@ fn test_snippet_unnamed_struct() {
     #[error("welp")]
     #[diagnostic(code(foo::bar::baz))]
     struct Foo(
-        String,
-        #[snippet(0, message("hi"))] SourceSpan,
-        #[highlight(1)] SourceSpan,
-        #[highlight(1)] SourceSpan,
-        // referenced source name
-        String,
-        #[snippet(0, message("{}", self.4))] SourceSpan,
-        #[highlight(5)] SourceSpan,
-        #[highlight(5)] SourceSpan,
+        #[source_code] String,
+        #[label("{0}")] SourceSpan,
+        #[label = "idk"] SourceSpan,
+        #[label] SourceSpan,
     );
 }
 
@@ -266,24 +235,23 @@ fn test_snippet_enum() {
     enum Foo {
         #[diagnostic(code(foo::a))]
         A {
+            #[source_code]
             src: String,
-            #[snippet(src, message("hi this is where the thing went wrong"))]
-            snip: SourceSpan,
-            #[highlight(snip)]
+            msg: String,
+            #[label("hi this is where the thing went wrong ({msg})")]
+            var0: SourceSpan,
+            #[label = "blorp"]
             var1: SourceSpan,
-            #[highlight(snip)]
+            #[label]
             var2: SourceSpan,
         },
         #[diagnostic(code(foo::b))]
         B(
+            #[source_code] String,
             String,
-            #[snippet(0, message("hi"))] SourceSpan,
-            #[highlight(1)] SourceSpan,
-            #[highlight(1, label("var 2"))] SourceSpan,
-            // referenced source name
-            #[snippet(0)] SourceSpan,
-            #[highlight(4)] SourceSpan,
-            #[highlight(4)] SourceSpan,
+            #[label("{1}")] SourceSpan,
+            #[label = "blorp"] SourceSpan,
+            #[label] SourceSpan,
         ),
     }
 }
@@ -329,19 +297,17 @@ const SNIPPET_TEXT: &str = "hello from miette";
     severity(Warning)
 )]
 struct ForwardsTo {
+    #[source_code]
     src: String,
-    #[snippet(src, message("snippet text"))]
-    snip: miette::SourceSpan,
-    #[highlight(snip, label("highlight text"))]
-    highlight: miette::SourceSpan,
+    #[label("highlight text")]
+    label: miette::SourceSpan,
 }
 
 impl ForwardsTo {
     fn new() -> Self {
         ForwardsTo {
             src: SNIPPET_TEXT.into(),
-            snip: SourceSpan::new(0.into(), SNIPPET_TEXT.len().into()),
-            highlight: SourceSpan::new(11.into(), 6.into()),
+            label: SourceSpan::new(11.into(), 6.into()),
         }
     }
 }
@@ -352,39 +318,6 @@ fn check_all(diag: &impl Diagnostic) {
     assert_eq!(diag.url().unwrap().to_string(), "https://example.com");
     assert_eq!(diag.help().unwrap().to_string(), "help");
     assert_eq!(diag.severity().unwrap(), miette::Severity::Warning);
-    check_snippets(diag);
-}
-
-fn check_snippets(diag: &impl Diagnostic) {
-    type Snip = (Option<String>, usize, usize);
-    let snips: Vec<(Snip, Vec<Snip>)> = diag
-        .snippets()
-        .unwrap()
-        .map(
-            |miette::DiagnosticSnippet {
-                 message,
-                 context,
-                 highlights,
-                 ..
-             }| {
-                (
-                    (message, context.offset(), context.len()),
-                    highlights
-                        .into_iter()
-                        .flatten()
-                        .map(|(msg, span)| (msg, span.offset(), span.len()))
-                        .collect(),
-                )
-            },
-        )
-        .collect();
-    assert_eq!(
-        &snips[..],
-        &[(
-            (Some("snippet text".into()), 0, SNIPPET_TEXT.len()),
-            vec![(Some("highlight text".into()), 11, 6)]
-        )]
-    );
 }
 
 #[test]
@@ -475,8 +408,6 @@ fn test_forward_struct_named() {
     assert_eq!(diag.code().unwrap().to_string(), "foo::bar::overridden");
     assert_eq!(diag.help().unwrap().to_string(), "overridden help please");
     assert_eq!(diag.severity(), Some(Severity::Advice));
-    // this comes from <ForwardsTo as Diagnostic>::snippets()
-    check_snippets(&diag);
 }
 
 #[test]
@@ -490,8 +421,6 @@ fn test_forward_struct_unnamed() {
     let diag = Struct(ForwardsTo::new(), "url here");
     assert_eq!(diag.code().unwrap().to_string(), "foo::bar::overridden");
     assert_eq!(diag.url().unwrap().to_string(), "url here");
-    // this comes from <ForwardsTo as Diagnostic>::snippets()
-    check_snippets(&diag);
 }
 
 #[test]
@@ -515,8 +444,6 @@ fn test_forward_enum_named() {
         variant.help().unwrap().to_string(),
         "overridden help please"
     );
-    // this comes from <ForwardsTo as Diagnostic>::snippets()
-    check_snippets(&variant);
 }
 
 #[test]
@@ -534,8 +461,6 @@ fn test_forward_enum_unnamed() {
         variant.help().unwrap().to_string(),
         "overridden help please"
     );
-    // this comes from <ForwardsTo as Diagnostic>::snippets()
-    check_snippets(&variant);
 }
 
 #[test]
diff --git a/tests/narrated.rs b/tests/narrated.rs
index 041f0f6..c9224ce 100644
--- a/tests/narrated.rs
+++ b/tests/narrated.rs
@@ -1,7 +1,4 @@
-use miette::{
-    Diagnostic, MietteError, NamedSource,
-    NarratableReportHandler, Report, SourceSpan,
-};
+use miette::{Diagnostic, MietteError, NamedSource, NarratableReportHandler, Report, SourceSpan};
 
 #[cfg(feature = "fancy")]
 use miette::{GraphicalReportHandler, GraphicalTheme};
@@ -30,19 +27,17 @@ fn single_line_highlight() -> Result<(), MietteError> {
     #[error("oops!")]
     #[diagnostic(code(oops::my::bad), help("try doing it better next time?"))]
     struct MyBad {
+        #[source_code]
         src: NamedSource,
-        #[snippet(src, message("This is the part that broke"))]
-        ctx: SourceSpan,
-        #[highlight(ctx, label = "this bit here")]
-        highlight: SourceSpan,
+        #[label("this bit here")]
+        bad_thing: SourceSpan,
     }
 
     let src = "source\n  text\n    here".to_string();
     let len = src.len();
     let err = MyBad {
         src: NamedSource::new("bad_file.rs", src),
-        ctx: (0, len).into(),
-        highlight: (9, 4).into(),
+        bad_thing: (9, 4).into(),
     };
     let out = fmt_report(err.into());
     println!("{}", out);