Support matching custom enums

4 files changed, 81 insertions, 7 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 701de02..356c927 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -1,6 +1,6 @@
 use filters;
 use input::TemplateInput;
-use parser::{self, Cond, Expr, Macro, MatchParameter, Node, Target, When, WS};
+use parser::{self, Cond, Expr, Macro, MatchParameter, MatchVariant, Node, Target, When, WS};
 use path;
 
 use quote::{Tokens, ToTokens};
@@ -635,19 +635,24 @@ impl<'a> Generator<'a> {
         }
     }
 
-    fn visit_match_variant(&mut self, param: &MatchParameter) -> DisplayWrap {
+    fn visit_match_variant(&mut self, param: &MatchVariant) -> DisplayWrap {
         match *param {
-            MatchParameter::StrLit(s) => self.visit_str_lit(s),
-            MatchParameter::NumLit(s) => {
+            MatchVariant::StrLit(s) => self.visit_str_lit(s),
+            MatchVariant::NumLit(s) => {
                 // Variants need to be references until match-modes land
                 self.write("&");
                 self.visit_num_lit(s)
             },
-            MatchParameter::Name(s) => {
+            MatchVariant::Name(s) => {
                 self.write("&");
                 self.write(s);
                 DisplayWrap::Unwrapped
             }
+            MatchVariant::Path(ref s) => {
+                self.write("&");
+                self.write(&s.join("::"));
+                DisplayWrap::Unwrapped
+            }
         }
     }
 
diff --git a/askama_shared/src/parser.rs b/askama_shared/src/parser.rs
index 76125ef..338bf1e 100644
--- a/askama_shared/src/parser.rs
+++ b/askama_shared/src/parser.rs
@@ -16,6 +16,14 @@ pub enum Expr<'a> {
 }
 
 #[derive(Debug)]
+pub enum MatchVariant<'a> {
+    Path(Vec<&'a str>),
+    Name(&'a str),
+    NumLit(&'a str),
+    StrLit(&'a str),
+}
+
+#[derive(Debug)]
 pub enum MatchParameter<'a> {
     Name(&'a str),
     NumLit(&'a str),
@@ -57,7 +65,7 @@ pub enum Node<'a> {
 }
 
 pub type Cond<'a> = (WS, Option<Expr<'a>>, Vec<Node<'a>>);
-pub type When<'a> = (WS, Option<MatchParameter<'a>>, Vec<MatchParameter<'a>>, Vec<Node<'a>>);
+pub type When<'a> = (WS, Option<MatchVariant<'a>>, Vec<MatchParameter<'a>>, Vec<Node<'a>>);
 
 fn split_ws_parts(s: &[u8]) -> Node {
     if s.is_empty() {
@@ -154,6 +162,10 @@ named!(expr_array_lit<Expr>, do_parse!(
     })
 ));
 
+named!(variant_num_lit<MatchVariant>, map!(num_lit,
+    |s| MatchVariant::NumLit(s)
+));
+
 named!(param_num_lit<MatchParameter>, map!(num_lit,
     |s| MatchParameter::NumLit(s)
 ));
@@ -163,6 +175,11 @@ named!(expr_str_lit<Expr>, map!(
     |s| Expr::StrLit(str::from_utf8(s).unwrap())
 ));
 
+named!(variant_str_lit<MatchVariant>, map!(
+    delimited!(char!('"'), is_not!("\""), char!('"')),
+    |s| MatchVariant::StrLit(str::from_utf8(s).unwrap())
+));
+
 named!(param_str_lit<MatchParameter>, map!(
     delimited!(char!('"'), is_not!("\""), char!('"')),
     |s| MatchParameter::StrLit(str::from_utf8(s).unwrap())
@@ -186,10 +203,28 @@ named!(expr_path<Expr>, do_parse!(
     })
 ));
 
+named!(variant_path<MatchVariant>, do_parse!(
+    start: call!(identifier) >>
+    rest: many1!(do_parse!(
+        tag_s!("::") >>
+        part: identifier >>
+        (part)
+    )) >>
+    ({
+        let mut path = vec![start];
+        path.extend(rest);
+        MatchVariant::Path(path)
+    })
+));
+
 named!(target_single<Target>, map!(identifier,
     |s| Target::Name(s)
 ));
 
+named!(variant_name<MatchVariant>, map!(identifier,
+    |s| MatchVariant::Name(s)
+));
+
 named!(param_name<MatchParameter>, map!(identifier,
     |s| MatchParameter::Name(s)
 ));
@@ -266,6 +301,13 @@ named!(expr_single<Expr>, alt!(
     expr_group
 ));
 
+named!(match_variant<MatchVariant>, alt!(
+    variant_path |
+    variant_name |
+    variant_num_lit |
+    variant_str_lit
+));
+
 named!(match_parameter<MatchParameter>, alt!(
     param_name |
     param_num_lit |
@@ -421,7 +463,7 @@ named!(when_block<When>, do_parse!(
     tag_s!("{%") >>
     pws: opt!(tag_s!("-")) >>
     ws!(tag_s!("when")) >>
-    variant: ws!(match_parameter) >>
+    variant: ws!(match_variant) >>
     params: opt!(ws!(with_parameters)) >>
     nws: opt!(tag_s!("-")) >>
     tag_s!("%}") >>
diff --git a/testing/templates/match-custom-enum.html b/testing/templates/match-custom-enum.html
new file mode 100644
index 0000000..cb45b8f
--- /dev/null
+++ b/testing/templates/match-custom-enum.html
@@ -0,0 +1,8 @@
+{% match color %}
+{% when Color::Rgb with (r, g, b) %}
+Colorful: #{{ "{:02X}"|format(r) }}{{ "{:02X}"|format(g) }}{{ "{:02X}"|format(b) }}
+{% when Color::GrayScale with (val) %}
+Gray: #{{ "{:02X}"|format(val) }}{{ "{:02X}"|format(val) }}{{ "{:02X}"|format(val) }}
+{% else %}
+CMYK not supported
+{% endmatch %}
diff --git a/testing/tests/matches.rs b/testing/tests/matches.rs
index b8a6c98..86d32be 100644
--- a/testing/tests/matches.rs
+++ b/testing/tests/matches.rs
@@ -63,3 +63,22 @@ fn test_match_literal_num() {
     let s = MatchLitNumTemplate { item: 23 };
     assert_eq!(s.render().unwrap(), "\n\nElse found 23\n");
 }
+
+#[allow(dead_code)]
+enum Color {
+    Rgb(u32, u32, u32),
+    GrayScale(u32),
+    Cmyk(u32, u32, u32, u32)
+}
+
+#[derive(Template)]
+#[template(path = "match-custom-enum.html")]
+struct MatchCustomEnumTemplate {
+    color: Color,
+}
+
+#[test]
+fn test_match_custom_enum() {
+    let s = MatchCustomEnumTemplate { color: Color::Rgb(160, 0, 255) };
+    assert_eq!(s.render().unwrap(), "\n\nColorful: #A000FF\n");
+}