Store extends path in Context

2 files changed, 14 insertions, 15 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 884ab3e..5d16f2c 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -58,7 +58,7 @@ impl<'a> Generator<'a> {
     // Takes a Context and generates the relevant implementations.
     fn build(mut self, ctx: &'a Context) -> String {
         if !ctx.blocks.is_empty() {
-            if !ctx.derived {
+            if ctx.extends.is_none() {
                 self.define_trait(ctx);
             } else {
                 let parent_type = get_parent_type(self.input.ast)
@@ -66,7 +66,7 @@ impl<'a> Generator<'a> {
                 self.deref_to_parent(parent_type);
             }
 
-            let trait_nodes = if !ctx.derived {
+            let trait_nodes = if ctx.extends.is_none() {
                 Some(&ctx.nodes[..])
             } else {
                 None
@@ -156,7 +156,7 @@ impl<'a> Generator<'a> {
             "fn render_into(&self, writer: &mut ::std::fmt::Write) \
              -> ::askama::Result<()> {",
         );
-        if ctx.derived {
+        if ctx.extends.is_some() {
             self.writeln("self._parent.render_trait_into(self, writer)");
         } else {
             self.writeln("self.render_trait_into(self, writer)");
diff --git a/askama_derive/src/lib.rs b/askama_derive/src/lib.rs
index 3651e02..07ba0a6 100644
--- a/askama_derive/src/lib.rs
+++ b/askama_derive/src/lib.rs
@@ -91,26 +91,26 @@ fn find_used_templates(map: &mut HashMap<PathBuf, String>, path: PathBuf, source
 
 pub(crate) struct Context<'a> {
     nodes: &'a [Node<'a>],
+    extends: Option<PathBuf>,
     blocks: Vec<&'a Node<'a>>,
     macros: HashMap<&'a str, &'a Macro<'a>>,
     imports: HashMap<&'a str, PathBuf>,
     trait_name: String,
-    derived: bool,
 }
 
 impl<'a> Context<'a> {
     fn new<'n>(input: &'n TemplateInput, nodes: &'n [Node<'n>]) -> Context<'n> {
-        let mut base = None;
+        let mut extends = None;
         let mut blocks = Vec::new();
         let mut macros = HashMap::new();
         let mut imports = HashMap::new();
 
         for n in nodes {
             match n {
-                Node::Extends(Expr::StrLit(path)) => match base {
+                Node::Extends(Expr::StrLit(path)) => match extends {
                     Some(_) => panic!("multiple extend blocks found"),
                     None => {
-                        base = Some(*path);
+                        extends = Some(path::find_template_from_path(path, Some(&input.path)));
                     }
                 },
                 def @ Node::BlockDef(_, _, _, _) => {
@@ -143,19 +143,18 @@ impl<'a> Context<'a> {
             check_nested += 1;
         }
 
+        let trait_name = match extends {
+            Some(ref path) => trait_name_for_path(path),
+            None => trait_name_for_path(&input.path),
+        };
+
         Context {
             nodes,
+            extends,
             blocks,
             macros,
             imports,
-            trait_name: match base {
-                Some(user_path) => trait_name_for_path(&path::find_template_from_path(
-                    user_path,
-                    Some(&input.path),
-                )),
-                None => trait_name_for_path(&input.path),
-            },
-            derived: base.is_some(),
+            trait_name,
         }
     }
 }