diff --git a/askama_derive/src/config.rs b/askama_derive/src/config.rs
index b2a13aa82..47801cdcb 100644
--- a/askama_derive/src/config.rs
+++ b/askama_derive/src/config.rs
@@ -6,7 +6,6 @@ use std::{env, fs};
 use serde::Deserialize;
 
 use crate::CompileError;
-
 use parser::node::Whitespace;
 use parser::Syntax;
 
diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 76b775b8b..7aa6d296d 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -1,250 +1,19 @@
-use crate::config::{get_template_source, read_config_file, Config, WhitespaceHandling};
+use std::collections::hash_map::{Entry, HashMap};
+use std::path::{Path, PathBuf};
+use std::{cmp, hash, mem, str};
+
+use crate::config::{get_template_source, WhitespaceHandling};
 use crate::heritage::{Context, Heritage};
-use crate::input::{Print, Source, TemplateInput};
+use crate::input::{Source, TemplateInput};
 use crate::CompileError;
 
 use parser::node::{
     Call, Comment, CondTest, If, Include, Let, Lit, Loop, Match, Target, Whitespace, Ws,
 };
 use parser::{Expr, Node, Parsed};
-use proc_macro::TokenStream;
-use quote::{quote, ToTokens};
-use syn::punctuated::Punctuated;
-
-use std::collections::hash_map::{Entry, HashMap};
-use std::path::{Path, PathBuf};
-use std::{cmp, hash, mem, str};
-
-/// The actual implementation for askama_derive::Template
-pub(crate) fn derive_template(input: TokenStream) -> TokenStream {
-    let ast: syn::DeriveInput = syn::parse(input).unwrap();
-    match build_template(&ast) {
-        Ok(source) => source.parse().unwrap(),
-        Err(e) => e.into_compile_error(),
-    }
-}
-
-/// Takes a `syn::DeriveInput` and generates source code for it
-///
-/// Reads the metadata from the `template()` attribute to get the template
-/// metadata, then fetches the source from the filesystem. The source is
-/// parsed, and the parse tree is fed to the code generator. Will print
-/// the parse tree and/or generated source according to the `print` key's
-/// value as passed to the `template()` attribute.
-fn build_template(ast: &syn::DeriveInput) -> Result<String, CompileError> {
-    let template_args = TemplateArgs::new(ast)?;
-    let config_toml = read_config_file(template_args.config_path.as_deref())?;
-    let config = Config::new(&config_toml, template_args.whitespace.as_ref())?;
-    let input = TemplateInput::new(ast, &config, template_args)?;
-    let source: String = match input.source {
-        Source::Source(ref s) => s.clone(),
-        Source::Path(_) => get_template_source(&input.path)?,
-    };
-
-    let mut templates = HashMap::new();
-    find_used_templates(&input, &mut templates, source)?;
-
-    let mut contexts = HashMap::new();
-    for (path, parsed) in &templates {
-        contexts.insert(
-            path.as_path(),
-            Context::new(input.config, path, parsed.nodes())?,
-        );
-    }
-
-    let ctx = &contexts[input.path.as_path()];
-    let heritage = if !ctx.blocks.is_empty() || ctx.extends.is_some() {
-        Some(Heritage::new(ctx, &contexts))
-    } else {
-        None
-    };
-
-    if input.print == Print::Ast || input.print == Print::All {
-        eprintln!("{:?}", templates[input.path.as_path()].nodes());
-    }
-
-    let code = Generator::new(
-        &input,
-        &contexts,
-        heritage.as_ref(),
-        MapChain::new(),
-        config.whitespace,
-    )
-    .build(&contexts[input.path.as_path()])?;
-    if input.print == Print::Code || input.print == Print::All {
-        eprintln!("{code}");
-    }
-    Ok(code)
-}
-
-#[derive(Default)]
-pub(crate) struct TemplateArgs {
-    pub(crate) source: Option<Source>,
-    pub(crate) print: Print,
-    pub(crate) escaping: Option<String>,
-    pub(crate) ext: Option<String>,
-    pub(crate) syntax: Option<String>,
-    pub(crate) config_path: Option<String>,
-    pub(crate) whitespace: Option<String>,
-}
-
-impl TemplateArgs {
-    fn new(ast: &'_ syn::DeriveInput) -> Result<Self, CompileError> {
-        // Check that an attribute called `template()` exists once and that it is
-        // the proper type (list).
-        let mut template_args = None;
-        for attr in &ast.attrs {
-            if !attr.path().is_ident("template") {
-                continue;
-            }
-
-            match attr.parse_args_with(Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated) {
-                Ok(args) if template_args.is_none() => template_args = Some(args),
-                Ok(_) => return Err("duplicated 'template' attribute".into()),
-                Err(e) => return Err(format!("unable to parse template arguments: {e}").into()),
-            };
-        }
-
-        let template_args =
-            template_args.ok_or_else(|| CompileError::from("no attribute 'template' found"))?;
-
-        let mut args = Self::default();
-        // Loop over the meta attributes and find everything that we
-        // understand. Return a CompileError if something is not right.
-        // `source` contains an enum that can represent `path` or `source`.
-        for item in template_args {
-            let pair = match item {
-                syn::Meta::NameValue(pair) => pair,
-                _ => {
-                    return Err(format!(
-                        "unsupported attribute argument {:?}",
-                        item.to_token_stream()
-                    )
-                    .into())
-                }
-            };
-
-            let ident = match pair.path.get_ident() {
-                Some(ident) => ident,
-                None => unreachable!("not possible in syn::Meta::NameValue(…)"),
-            };
+use quote::quote;
 
-            let value = match pair.value {
-                syn::Expr::Lit(lit) => lit,
-                syn::Expr::Group(group) => match *group.expr {
-                    syn::Expr::Lit(lit) => lit,
-                    _ => {
-                        return Err(format!("unsupported argument value type for {ident:?}").into())
-                    }
-                },
-                _ => return Err(format!("unsupported argument value type for {ident:?}").into()),
-            };
-
-            if ident == "path" {
-                if let syn::Lit::Str(s) = value.lit {
-                    if args.source.is_some() {
-                        return Err("must specify 'source' or 'path', not both".into());
-                    }
-                    args.source = Some(Source::Path(s.value()));
-                } else {
-                    return Err("template path must be string literal".into());
-                }
-            } else if ident == "source" {
-                if let syn::Lit::Str(s) = value.lit {
-                    if args.source.is_some() {
-                        return Err("must specify 'source' or 'path', not both".into());
-                    }
-                    args.source = Some(Source::Source(s.value()));
-                } else {
-                    return Err("template source must be string literal".into());
-                }
-            } else if ident == "print" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.print = s.value().parse()?;
-                } else {
-                    return Err("print value must be string literal".into());
-                }
-            } else if ident == "escape" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.escaping = Some(s.value());
-                } else {
-                    return Err("escape value must be string literal".into());
-                }
-            } else if ident == "ext" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.ext = Some(s.value());
-                } else {
-                    return Err("ext value must be string literal".into());
-                }
-            } else if ident == "syntax" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.syntax = Some(s.value())
-                } else {
-                    return Err("syntax value must be string literal".into());
-                }
-            } else if ident == "config" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.config_path = Some(s.value())
-                } else {
-                    return Err("config value must be string literal".into());
-                }
-            } else if ident == "whitespace" {
-                if let syn::Lit::Str(s) = value.lit {
-                    args.whitespace = Some(s.value())
-                } else {
-                    return Err("whitespace value must be string literal".into());
-                }
-            } else {
-                return Err(format!("unsupported attribute key {ident:?} found").into());
-            }
-        }
-
-        Ok(args)
-    }
-}
-
-fn find_used_templates(
-    input: &TemplateInput<'_>,
-    map: &mut HashMap<PathBuf, Parsed>,
-    source: String,
-) -> Result<(), CompileError> {
-    let mut dependency_graph = Vec::new();
-    let mut check = vec![(input.path.clone(), source)];
-    while let Some((path, source)) = check.pop() {
-        let parsed = Parsed::new(source, input.syntax)?;
-        for n in parsed.nodes() {
-            match n {
-                Node::Extends(extends) => {
-                    let extends = input.config.find_template(extends.path, Some(&path))?;
-                    let dependency_path = (path.clone(), extends.clone());
-                    if dependency_graph.contains(&dependency_path) {
-                        return Err(format!(
-                            "cyclic dependency in graph {:#?}",
-                            dependency_graph
-                                .iter()
-                                .map(|e| format!("{:#?} --> {:#?}", e.0, e.1))
-                                .collect::<Vec<String>>()
-                        )
-                        .into());
-                    }
-                    dependency_graph.push(dependency_path);
-                    let source = get_template_source(&extends)?;
-                    check.push((extends, source));
-                }
-                Node::Import(import) => {
-                    let import = input.config.find_template(import.path, Some(&path))?;
-                    let source = get_template_source(&import)?;
-                    check.push((import, source));
-                }
-                _ => {}
-            }
-        }
-        map.insert(path, parsed);
-    }
-    Ok(())
-}
-
-struct Generator<'a> {
+pub(crate) struct Generator<'a> {
     // The template input state: original struct AST and attributes
     input: &'a TemplateInput<'a>,
     // All contexts, keyed by the package-relative template path
@@ -268,18 +37,14 @@ struct Generator<'a> {
     buf_writable: Vec<Writable<'a>>,
     // Counter for write! hash named arguments
     named: usize,
-    // If set to `suppress`, the whitespace characters will be removed by default unless `+` is
-    // used.
-    whitespace: WhitespaceHandling,
 }
 
 impl<'a> Generator<'a> {
-    fn new<'n>(
+    pub(crate) fn new<'n>(
         input: &'n TemplateInput<'_>,
         contexts: &'n HashMap<&'n Path, Context<'n>>,
         heritage: Option<&'n Heritage<'_>>,
         locals: MapChain<'n, &'n str, LocalMeta>,
-        whitespace: WhitespaceHandling,
     ) -> Generator<'n> {
         Generator {
             input,
@@ -292,12 +57,11 @@ impl<'a> Generator<'a> {
             super_block: None,
             buf_writable: vec![],
             named: 0,
-            whitespace,
         }
     }
 
     // Takes a Context and generates the relevant implementations.
-    fn build(mut self, ctx: &'a Context<'_>) -> Result<String, CompileError> {
+    pub(crate) fn build(mut self, ctx: &'a Context<'_>) -> Result<String, CompileError> {
         let mut buf = Buffer::new(0);
 
         self.impl_template(ctx, &mut buf)?;
@@ -1027,13 +791,7 @@ impl<'a> Generator<'a> {
         // handle the include's nodes. Unfortunately we can't easily share the `includes` cache.
 
         let locals = MapChain::with_parent(&self.locals);
-        let mut child = Self::new(
-            self.input,
-            self.contexts,
-            self.heritage,
-            locals,
-            self.whitespace,
-        );
+        let mut child = Self::new(self.input, self.contexts, self.heritage, locals);
 
         let nodes = match self.contexts.get(path.as_path()) {
             Some(ctx) => ctx.nodes,
@@ -1869,7 +1627,7 @@ impl<'a> Generator<'a> {
             Some(Whitespace::Suppress) => WhitespaceHandling::Suppress,
             Some(Whitespace::Preserve) => WhitespaceHandling::Preserve,
             Some(Whitespace::Minimize) => WhitespaceHandling::Minimize,
-            None => self.whitespace,
+            None => self.input.config.whitespace,
         }
     }
 
@@ -1970,7 +1728,7 @@ impl Buffer {
 }
 
 #[derive(Clone, Default)]
-struct LocalMeta {
+pub(crate) struct LocalMeta {
     refs: Option<String>,
     initialized: bool,
 }
@@ -1994,7 +1752,7 @@ impl LocalMeta {
 // type SetChain<'a, T> = MapChain<'a, T, ()>;
 
 #[derive(Debug)]
-struct MapChain<'a, K, V>
+pub(crate) struct MapChain<'a, K, V>
 where
     K: cmp::Eq + hash::Hash,
 {
@@ -2006,13 +1764,6 @@ impl<'a, K: 'a, V: 'a> MapChain<'a, K, V>
 where
     K: cmp::Eq + hash::Hash,
 {
-    fn new() -> MapChain<'a, K, V> {
-        MapChain {
-            parent: None,
-            scopes: vec![HashMap::new()],
-        }
-    }
-
     fn with_parent<'p>(parent: &'p MapChain<'_, K, V>) -> MapChain<'p, K, V> {
         MapChain {
             parent: Some(parent),
@@ -2075,6 +1826,15 @@ impl MapChain<'_, &str, LocalMeta> {
     }
 }
 
+impl<'a, K: Eq + hash::Hash, V> Default for MapChain<'a, K, V> {
+    fn default() -> Self {
+        Self {
+            parent: None,
+            scopes: vec![HashMap::new()],
+        }
+    }
+}
+
 /// Returns `true` if enough assumptions can be made,
 /// to determine that `self` is copyable.
 fn is_copyable(expr: &Expr<'_>) -> bool {
diff --git a/askama_derive/src/heritage.rs b/askama_derive/src/heritage.rs
index ce85ac769..d75d0a597 100644
--- a/askama_derive/src/heritage.rs
+++ b/askama_derive/src/heritage.rs
@@ -3,7 +3,6 @@ use std::path::{Path, PathBuf};
 
 use crate::config::Config;
 use crate::CompileError;
-
 use parser::node::{BlockDef, Macro, Match};
 use parser::Node;
 
@@ -26,7 +25,7 @@ impl Heritage<'_> {
         while let Some(ref path) = ctx.extends {
             ctx = &contexts[path.as_path()];
             for (name, def) in &ctx.blocks {
-                blocks.entry(name).or_insert_with(Vec::new).push((ctx, def));
+                blocks.entry(name).or_default().push((ctx, def));
             }
         }
 
diff --git a/askama_derive/src/input.rs b/askama_derive/src/input.rs
index d4e8ad920..ee9b94233 100644
--- a/askama_derive/src/input.rs
+++ b/askama_derive/src/input.rs
@@ -1,21 +1,23 @@
-use crate::config::Config;
-use crate::generator::TemplateArgs;
-use crate::CompileError;
-use parser::Syntax;
-
+use std::collections::hash_map::HashMap;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
 
 use mime::Mime;
+use quote::ToTokens;
+use syn::punctuated::Punctuated;
+
+use crate::config::{get_template_source, read_config_file, Config};
+use crate::CompileError;
+use parser::{Node, Parsed, Syntax};
 
 pub(crate) struct TemplateInput<'a> {
     pub(crate) ast: &'a syn::DeriveInput,
     pub(crate) config: &'a Config<'a>,
     pub(crate) syntax: &'a Syntax<'a>,
-    pub(crate) source: Source,
+    pub(crate) source: &'a Source,
     pub(crate) print: Print,
     pub(crate) escaper: &'a str,
-    pub(crate) ext: Option<String>,
+    pub(crate) ext: Option<&'a str>,
     pub(crate) mime_type: String,
     pub(crate) path: PathBuf,
 }
@@ -27,7 +29,7 @@ impl TemplateInput<'_> {
     pub(crate) fn new<'n>(
         ast: &'n syn::DeriveInput,
         config: &'n Config<'_>,
-        args: TemplateArgs,
+        args: &'n TemplateArgs,
     ) -> Result<TemplateInput<'n>, CompileError> {
         let TemplateArgs {
             source,
@@ -41,7 +43,9 @@ impl TemplateInput<'_> {
         // Validate the `source` and `ext` value together, since they are
         // related. In case `source` was used instead of `path`, the value
         // of `ext` is merged into a synthetic `path` value here.
-        let source = source.expect("template path or source not found in attributes");
+        let source = source
+            .as_ref()
+            .expect("template path or source not found in attributes");
         let path = match (&source, &ext) {
             (Source::Path(path), _) => config.find_template(path, None)?,
             (&Source::Source(_), Some(ext)) => PathBuf::from(format!("{}.{}", ast.ident, ext)),
@@ -51,28 +55,25 @@ impl TemplateInput<'_> {
         };
 
         // Validate syntax
-        let syntax = syntax.map_or_else(
+        let syntax = syntax.as_deref().map_or_else(
             || Ok(config.syntaxes.get(config.default_syntax).unwrap()),
             |s| {
                 config
                     .syntaxes
-                    .get(&s)
+                    .get(s)
                     .ok_or_else(|| CompileError::from(format!("attribute syntax {s} not exist")))
             },
         )?;
 
         // Match extension against defined output formats
 
-        let escaping = escaping.unwrap_or_else(|| {
-            path.extension()
-                .map(|s| s.to_str().unwrap())
-                .unwrap_or("")
-                .to_string()
-        });
+        let escaping = escaping
+            .as_deref()
+            .unwrap_or_else(|| path.extension().map(|s| s.to_str().unwrap()).unwrap_or(""));
 
         let mut escaper = None;
         for (extensions, path) in &config.escapers {
-            if extensions.contains(&escaping) {
+            if extensions.contains(escaping) {
                 escaper = Some(path);
                 break;
             }
@@ -91,17 +92,192 @@ impl TemplateInput<'_> {
             config,
             syntax,
             source,
-            print,
+            print: *print,
             escaper,
-            ext,
+            ext: ext.as_deref(),
             mime_type,
             path,
         })
     }
 
+    pub(crate) fn find_used_templates(
+        &self,
+        map: &mut HashMap<PathBuf, Parsed>,
+    ) -> Result<(), CompileError> {
+        let source = match &self.source {
+            Source::Source(s) => s.clone(),
+            Source::Path(_) => get_template_source(&self.path)?,
+        };
+
+        let mut dependency_graph = Vec::new();
+        let mut check = vec![(self.path.clone(), source)];
+        while let Some((path, source)) = check.pop() {
+            let parsed = Parsed::new(source, self.syntax)?;
+            for n in parsed.nodes() {
+                match n {
+                    Node::Extends(extends) => {
+                        let extends = self.config.find_template(extends.path, Some(&path))?;
+                        let dependency_path = (path.clone(), extends.clone());
+                        if dependency_graph.contains(&dependency_path) {
+                            return Err(format!(
+                                "cyclic dependency in graph {:#?}",
+                                dependency_graph
+                                    .iter()
+                                    .map(|e| format!("{:#?} --> {:#?}", e.0, e.1))
+                                    .collect::<Vec<String>>()
+                            )
+                            .into());
+                        }
+                        dependency_graph.push(dependency_path);
+                        let source = get_template_source(&extends)?;
+                        check.push((extends, source));
+                    }
+                    Node::Import(import) => {
+                        let import = self.config.find_template(import.path, Some(&path))?;
+                        let source = get_template_source(&import)?;
+                        check.push((import, source));
+                    }
+                    _ => {}
+                }
+            }
+            map.insert(path, parsed);
+        }
+        Ok(())
+    }
+
     #[inline]
     pub(crate) fn extension(&self) -> Option<&str> {
-        ext_default_to_path(self.ext.as_deref(), &self.path)
+        ext_default_to_path(self.ext, &self.path)
+    }
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct TemplateArgs {
+    source: Option<Source>,
+    print: Print,
+    escaping: Option<String>,
+    ext: Option<String>,
+    syntax: Option<String>,
+    config: String,
+    whitespace: Option<String>,
+}
+
+impl TemplateArgs {
+    pub(crate) fn new(ast: &'_ syn::DeriveInput) -> Result<Self, CompileError> {
+        // Check that an attribute called `template()` exists once and that it is
+        // the proper type (list).
+        let mut template_args = None;
+        for attr in &ast.attrs {
+            if !attr.path().is_ident("template") {
+                continue;
+            }
+
+            match attr.parse_args_with(Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated) {
+                Ok(args) if template_args.is_none() => template_args = Some(args),
+                Ok(_) => return Err("duplicated 'template' attribute".into()),
+                Err(e) => return Err(format!("unable to parse template arguments: {e}").into()),
+            };
+        }
+
+        let template_args =
+            template_args.ok_or_else(|| CompileError::from("no attribute 'template' found"))?;
+
+        let mut args = Self::default();
+        // Loop over the meta attributes and find everything that we
+        // understand. Return a CompileError if something is not right.
+        // `source` contains an enum that can represent `path` or `source`.
+        for item in template_args {
+            let pair = match item {
+                syn::Meta::NameValue(pair) => pair,
+                _ => {
+                    return Err(format!(
+                        "unsupported attribute argument {:?}",
+                        item.to_token_stream()
+                    )
+                    .into())
+                }
+            };
+
+            let ident = match pair.path.get_ident() {
+                Some(ident) => ident,
+                None => unreachable!("not possible in syn::Meta::NameValue(…)"),
+            };
+
+            let value = match pair.value {
+                syn::Expr::Lit(lit) => lit,
+                syn::Expr::Group(group) => match *group.expr {
+                    syn::Expr::Lit(lit) => lit,
+                    _ => {
+                        return Err(format!("unsupported argument value type for {ident:?}").into())
+                    }
+                },
+                _ => return Err(format!("unsupported argument value type for {ident:?}").into()),
+            };
+
+            if ident == "path" {
+                if let syn::Lit::Str(s) = value.lit {
+                    if args.source.is_some() {
+                        return Err("must specify 'source' or 'path', not both".into());
+                    }
+                    args.source = Some(Source::Path(s.value()));
+                } else {
+                    return Err("template path must be string literal".into());
+                }
+            } else if ident == "source" {
+                if let syn::Lit::Str(s) = value.lit {
+                    if args.source.is_some() {
+                        return Err("must specify 'source' or 'path', not both".into());
+                    }
+                    args.source = Some(Source::Source(s.value()));
+                } else {
+                    return Err("template source must be string literal".into());
+                }
+            } else if ident == "print" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.print = s.value().parse()?;
+                } else {
+                    return Err("print value must be string literal".into());
+                }
+            } else if ident == "escape" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.escaping = Some(s.value());
+                } else {
+                    return Err("escape value must be string literal".into());
+                }
+            } else if ident == "ext" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.ext = Some(s.value());
+                } else {
+                    return Err("ext value must be string literal".into());
+                }
+            } else if ident == "syntax" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.syntax = Some(s.value())
+                } else {
+                    return Err("syntax value must be string literal".into());
+                }
+            } else if ident == "config" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.config = read_config_file(Some(&s.value()))?;
+                } else {
+                    return Err("config value must be string literal".into());
+                }
+            } else if ident == "whitespace" {
+                if let syn::Lit::Str(s) = value.lit {
+                    args.whitespace = Some(s.value())
+                } else {
+                    return Err("whitespace value must be string literal".into());
+                }
+            } else {
+                return Err(format!("unsupported attribute key {ident:?} found").into());
+            }
+        }
+
+        Ok(args)
+    }
+
+    pub(crate) fn config(&self) -> Result<Config<'_>, CompileError> {
+        Config::new(&self.config, self.whitespace.as_ref())
     }
 }
 
@@ -124,12 +300,13 @@ fn extension(path: &Path) -> Option<&str> {
     }
 }
 
+#[derive(Debug)]
 pub(crate) enum Source {
     Path(String),
     Source(String),
 }
 
-#[derive(PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub(crate) enum Print {
     All,
     Ast,
diff --git a/askama_derive/src/lib.rs b/askama_derive/src/lib.rs
index 8a737aa9a..547a449f3 100644
--- a/askama_derive/src/lib.rs
+++ b/askama_derive/src/lib.rs
@@ -1,8 +1,8 @@
 #![deny(elided_lifetimes_in_paths)]
 #![deny(unreachable_pub)]
 
-use std::borrow::Cow;
 use std::fmt;
+use std::{borrow::Cow, collections::HashMap};
 
 use proc_macro::TokenStream;
 use proc_macro2::Span;
@@ -11,12 +11,61 @@ use parser::ParseError;
 
 mod config;
 mod generator;
+use generator::{Generator, MapChain};
 mod heritage;
+use heritage::{Context, Heritage};
 mod input;
+use input::{Print, TemplateArgs, TemplateInput};
 
 #[proc_macro_derive(Template, attributes(template))]
 pub fn derive_template(input: TokenStream) -> TokenStream {
-    generator::derive_template(input)
+    let ast = syn::parse::<syn::DeriveInput>(input).unwrap();
+    match build_template(&ast) {
+        Ok(source) => source.parse().unwrap(),
+        Err(e) => e.into_compile_error(),
+    }
+}
+
+/// Takes a `syn::DeriveInput` and generates source code for it
+///
+/// Reads the metadata from the `template()` attribute to get the template
+/// metadata, then fetches the source from the filesystem. The source is
+/// parsed, and the parse tree is fed to the code generator. Will print
+/// the parse tree and/or generated source according to the `print` key's
+/// value as passed to the `template()` attribute.
+pub(crate) fn build_template(ast: &syn::DeriveInput) -> Result<String, CompileError> {
+    let template_args = TemplateArgs::new(ast)?;
+    let config = template_args.config()?;
+    let input = TemplateInput::new(ast, &config, &template_args)?;
+
+    let mut templates = HashMap::new();
+    input.find_used_templates(&mut templates)?;
+
+    let mut contexts = HashMap::new();
+    for (path, parsed) in &templates {
+        contexts.insert(
+            path.as_path(),
+            Context::new(input.config, path, parsed.nodes())?,
+        );
+    }
+
+    let ctx = &contexts[input.path.as_path()];
+    let heritage = if !ctx.blocks.is_empty() || ctx.extends.is_some() {
+        Some(Heritage::new(ctx, &contexts))
+    } else {
+        None
+    };
+
+    if input.print == Print::Ast || input.print == Print::All {
+        eprintln!("{:?}", templates[input.path.as_path()].nodes());
+    }
+
+    let code = Generator::new(&input, &contexts, heritage.as_ref(), MapChain::default())
+        .build(&contexts[input.path.as_path()])?;
+    if input.print == Print::Code || input.print == Print::All {
+        eprintln!("{code}");
+    }
+    Ok(code)
 }
 
 #[derive(Debug, Clone)]