Add possibility to set up highlighten prefix and postfix for snippet

src/snippet/mod.rs

@@ -11,6 +11,9 @@ use crate::{Document, Score, Searcher};
 
 const DEFAULT_MAX_NUM_CHARS: usize = 150;
 
+const DEFAULT_HIGHLIGHTING_PREFIX: &str = "<b>";
+const DEFAULT_HIGHLIGHTING_POSTFIX: &str = "</b>";
+
 #[derive(Debug)]
 pub struct FragmentCandidate {
     score: Score,
@@ -50,35 +53,51 @@ impl FragmentCandidate {
 }
 
 /// `Snippet`
-/// Contains a fragment of a document, and some highlighed parts inside it.
+/// Contains a fragment of a document, and some highlighted parts inside it.
 #[derive(Debug)]
 pub struct Snippet {
     fragment: String,
     highlighted: Vec<Range<usize>>,
+    highlighting_prefix: String,
+    highlighting_postfix: String,
 }
 
-const HIGHLIGHTEN_PREFIX: &str = "<b>";
-const HIGHLIGHTEN_POSTFIX: &str = "</b>";
-
 impl Snippet {
-    /// Create a new, empty, `Snippet`
+    /// Create a new `Snippet`.
+    fn new(
+        fragment: &str,
+        highlighted: Vec<Range<usize>>,
+        highlighting_prefix: &str,
+        highlighting_postfix: &str,
+    ) -> Self {
+        Self {
+            fragment: fragment.to_string(),
+            highlighted,
+            highlighting_prefix: highlighting_prefix.to_string(),
+            highlighting_postfix: highlighting_postfix.to_string(),
+        }
+    }
+
+    /// Create a new, empty, `Snippet`.
     pub fn empty() -> Snippet {
         Snippet {
             fragment: String::new(),
             highlighted: Vec::new(),
+            highlighting_prefix: String::new(),
+            highlighting_postfix: String::new(),
         }
     }
 
-    /// Returns a hignlightned html from the `Snippet`.
+    /// Returns a highlighted html from the `Snippet`.
     pub fn to_html(&self) -> String {
         let mut html = String::new();
         let mut start_from: usize = 0;
 
         for item in self.highlighted.iter() {
             html.push_str(&encode_minimal(&self.fragment[start_from..item.start]));
-            html.push_str(HIGHLIGHTEN_PREFIX);
+            html.push_str(&self.highlighting_prefix);
             html.push_str(&encode_minimal(&self.fragment[item.clone()]));
-            html.push_str(HIGHLIGHTEN_POSTFIX);
+            html.push_str(&self.highlighting_postfix);
             start_from = item.end;
         }
         html.push_str(&encode_minimal(
@@ -92,7 +111,7 @@ impl Snippet {
         &self.fragment
     }
 
-    /// Returns a list of higlighted positions from the `Snippet`.
+    /// Returns a list of highlighted positions from the `Snippet`.
     pub fn highlighted(&self) -> &[Range<usize>] {
         &self.highlighted
     }
@@ -148,7 +167,12 @@ fn search_fragments<'a>(
 ///
 /// Takes a vector of `FragmentCandidate`s and the text.
 /// Figures out the best fragment from it and creates a snippet.
-fn select_best_fragment_combination(fragments: &[FragmentCandidate], text: &str) -> Snippet {
+fn select_best_fragment_combination(
+    fragments: &[FragmentCandidate],
+    text: &str,
+    highlighting_prefix: &str,
+    highlighting_postfix: &str,
+) -> Snippet {
     let best_fragment_opt = fragments.iter().max_by(|left, right| {
         let cmp_score = left
             .score
@@ -167,17 +191,16 @@ fn select_best_fragment_combination(fragments: &[FragmentCandidate], text: &str)
             .iter()
             .map(|item| item.start - fragment.start_offset..item.end - fragment.start_offset)
             .collect();
-        Snippet {
-            fragment: fragment_text.to_string(),
+        Snippet::new(
+            fragment_text,
             highlighted,
-        }
+            highlighting_prefix,
+            highlighting_postfix,
+        )
     } else {
-        // when there no fragments to chose from,
-        // for now create a empty snippet
-        Snippet {
-            fragment: String::new(),
-            highlighted: vec![],
-        }
+        // When there are no fragments to chose from,
+        // for now create an empty snippet.
+        Snippet::empty()
     }
 }
 
@@ -227,6 +250,8 @@ pub struct SnippetGenerator {
     tokenizer: TextAnalyzer,
     field: Field,
     max_num_chars: usize,
+    highlighting_prefix: String,
+    highlighting_postfix: String,
 }
 
 impl SnippetGenerator {
@@ -260,9 +285,17 @@ impl SnippetGenerator {
             tokenizer,
             field,
             max_num_chars: DEFAULT_MAX_NUM_CHARS,
+            highlighting_prefix: DEFAULT_HIGHLIGHTING_PREFIX.to_string(),
+            highlighting_postfix: DEFAULT_HIGHLIGHTING_POSTFIX.to_string(),
         })
     }
 
+    /// Sets highlighted prefix and postfix.
+    pub fn set_highlighted_elements(&mut self, prefix: &str, postfix: &str) {
+        self.highlighting_prefix = prefix.to_string();
+        self.highlighting_postfix = postfix.to_string()
+    }
+
     /// Sets a maximum number of chars.
     pub fn set_max_num_chars(&mut self, max_num_chars: usize) {
         self.max_num_chars = max_num_chars;
@@ -290,7 +323,12 @@ impl SnippetGenerator {
     pub fn snippet(&self, text: &str) -> Snippet {
         let fragment_candidates =
             search_fragments(&self.tokenizer, text, &self.terms_text, self.max_num_chars);
-        select_best_fragment_combination(&fragment_candidates[..], text)
+        select_best_fragment_combination(
+            &fragment_candidates[..],
+            text,
+            &self.highlighting_prefix,
+            &self.highlighting_postfix,
+        )
     }
 }
 
@@ -300,7 +338,10 @@ mod tests {
 
     use maplit::btreemap;
 
-    use super::{search_fragments, select_best_fragment_combination};
+    use super::{
+        search_fragments, select_best_fragment_combination, DEFAULT_HIGHLIGHTING_POSTFIX,
+        DEFAULT_HIGHLIGHTING_PREFIX,
+    };
     use crate::query::QueryParser;
     use crate::schema::{IndexRecordOption, Schema, TextFieldIndexing, TextOptions, TEXT};
     use crate::tokenizer::SimpleTokenizer;
@@ -333,15 +374,20 @@ Survey in 2016, 2017, and 2018."#;
             assert_eq!(first.score, 1.9);
             assert_eq!(first.stop_offset, 89);
         }
-        let snippet = select_best_fragment_combination(&fragments[..], TEST_TEXT);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            TEST_TEXT,
+            "<b class=\"super\">",
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(
             snippet.fragment,
             "Rust is a systems programming language sponsored by\nMozilla which describes it as a \
              \"safe"
         );
         assert_eq!(
             snippet.to_html(),
-            "<b>Rust</b> is a systems programming <b>language</b> sponsored by\nMozilla which \
+            "<b class=\"super\">Rust</b> is a systems programming <b class=\"super\">language</b> sponsored by\nMozilla which \
              describes it as a &quot;safe"
         )
     }
@@ -359,7 +405,12 @@ Survey in 2016, 2017, and 2018."#;
                 assert_eq!(first.score, 1.0);
                 assert_eq!(first.stop_offset, 17);
             }
-            let snippet = select_best_fragment_combination(&fragments[..], TEST_TEXT);
+            let snippet = select_best_fragment_combination(
+                &fragments[..],
+                TEST_TEXT,
+                DEFAULT_HIGHLIGHTING_PREFIX,
+                DEFAULT_HIGHLIGHTING_POSTFIX,
+            );
             assert_eq!(snippet.to_html(), "<b>Rust</b> is a systems")
         }
         {
@@ -374,7 +425,12 @@ Survey in 2016, 2017, and 2018."#;
                 assert_eq!(first.score, 0.9);
                 assert_eq!(first.stop_offset, 17);
             }
-            let snippet = select_best_fragment_combination(&fragments[..], TEST_TEXT);
+            let snippet = select_best_fragment_combination(
+                &fragments[..],
+                TEST_TEXT,
+                DEFAULT_HIGHLIGHTING_PREFIX,
+                DEFAULT_HIGHLIGHTING_POSTFIX,
+            );
             assert_eq!(snippet.to_html(), "programming <b>language</b>")
         }
     }
@@ -396,7 +452,12 @@ Survey in 2016, 2017, and 2018."#;
             assert_eq!(first.stop_offset, 7);
         }
 
-        let snippet = select_best_fragment_combination(&fragments[..], text);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            text,
+            DEFAULT_HIGHLIGHTING_PREFIX,
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(snippet.fragment, "c d");
         assert_eq!(snippet.to_html(), "<b>c</b> d");
     }
@@ -418,7 +479,12 @@ Survey in 2016, 2017, and 2018."#;
             assert_eq!(first.start_offset, 8);
         }
 
-        let snippet = select_best_fragment_combination(&fragments[..], text);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            text,
+            DEFAULT_HIGHLIGHTING_PREFIX,
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(snippet.fragment, "e f");
         assert_eq!(snippet.to_html(), "e <b>f</b>");
     }
@@ -441,7 +507,12 @@ Survey in 2016, 2017, and 2018."#;
             assert_eq!(first.start_offset, 0);
         }
 
-        let snippet = select_best_fragment_combination(&fragments[..], text);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            text,
+            DEFAULT_HIGHLIGHTING_PREFIX,
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(snippet.fragment, "e f g");
         assert_eq!(snippet.to_html(), "e <b>f</b> g");
     }
@@ -457,7 +528,12 @@ Survey in 2016, 2017, and 2018."#;
 
         assert_eq!(fragments.len(), 0);
 
-        let snippet = select_best_fragment_combination(&fragments[..], text);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            text,
+            DEFAULT_HIGHLIGHTING_PREFIX,
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(snippet.fragment, "");
         assert_eq!(snippet.to_html(), "");
     }
@@ -470,7 +546,12 @@ Survey in 2016, 2017, and 2018."#;
         let fragments = search_fragments(&From::from(SimpleTokenizer), text, &terms, 3);
         assert_eq!(fragments.len(), 0);
 
-        let snippet = select_best_fragment_combination(&fragments[..], text);
+        let snippet = select_best_fragment_combination(
+            &fragments[..],
+            text,
+            DEFAULT_HIGHLIGHTING_PREFIX,
+            DEFAULT_HIGHLIGHTING_POSTFIX,
+        );
         assert_eq!(snippet.fragment, "");
         assert_eq!(snippet.to_html(), "");
     }
@@ -566,4 +647,49 @@ Survey in 2016, 2017, and 2018."#;
         }
         Ok(())
     }
+
+    #[test]
+    fn test_snippet_generator_custom_highlighted_elements() -> crate::Result<()> {
+        let mut schema_builder = Schema::builder();
+        let text_options = TextOptions::default().set_indexing_options(
+            TextFieldIndexing::default()
+                .set_tokenizer("en_stem")
+                .set_index_option(IndexRecordOption::Basic),
+        );
+        let text_field = schema_builder.add_text_field("text", text_options);
+        let schema = schema_builder.build();
+        let index = Index::create_in_ram(schema);
+        {
+            // writing the segment
+            let mut index_writer = index.writer_for_tests()?;
+            let doc = doc!(text_field => TEST_TEXT);
+            index_writer.add_document(doc)?;
+            index_writer.commit()?;
+        }
+        let searcher = index.reader().unwrap().searcher();
+        let query_parser = QueryParser::for_index(&index, vec![text_field]);
+        let query = query_parser.parse_query("rust design").unwrap();
+        let mut snippet_generator =
+            SnippetGenerator::create(&searcher, &*query, text_field).unwrap();
+        {
+            let snippet = snippet_generator.snippet(TEST_TEXT);
+            assert_eq!(
+                snippet.to_html(),
+                "imperative-procedural paradigms. <b>Rust</b> is syntactically similar to \
+                 C++[according to whom?],\nbut its <b>designers</b> intend it to provide better \
+                 memory safety"
+            );
+        }
+        {
+            snippet_generator.set_max_num_chars(90);
+            snippet_generator.set_highlighted_elements("<q class=\"super\">", "</q>");
+            let snippet = snippet_generator.snippet(TEST_TEXT);
+            assert_eq!(
+                snippet.to_html(),
+                "<q class=\"super\">Rust</q> is syntactically similar to C++[according to whom?],\nbut its \
+                <q class=\"super\">designers</q> intend it to"
+            );
+        }
+        Ok(())
+    }
 }