text_splitter/splitter/
markdown.rs

1/*!
2# [`MarkdownSplitter`]
3Semantic splitting of Markdown documents. Tries to use as many semantic units from Markdown
4as possible, according to the Common Mark specification.
5*/
6
7use std::{iter::once, ops::Range};
8
9use either::Either;
10use itertools::Itertools;
11use pulldown_cmark::{Event, Options, Parser, Tag};
12
13use crate::{
14    splitter::{SemanticLevel, Splitter},
15    trim::Trim,
16    ChunkConfig, ChunkSizer, TryChunkSizer,
17};
18
19use super::ChunkCharIndex;
20
21/// Markdown splitter. Recursively splits chunks into the largest
22/// semantic units that fit within the chunk size. Also will
23/// attempt to merge neighboring chunks if they can fit within the
24/// given chunk size.
25#[derive(Debug)]
26pub struct MarkdownSplitter<Sizer>
27where
28    Sizer: TryChunkSizer,
29{
30    /// Method of determining chunk sizes.
31    chunk_config: ChunkConfig<Sizer>,
32}
33
34impl<Sizer> MarkdownSplitter<Sizer>
35where
36    Sizer: TryChunkSizer,
37{
38    /// Creates a new [`MarkdownSplitter`].
39    ///
40    /// ```
41    /// use text_splitter::MarkdownSplitter;
42    ///
43    /// // By default, the chunk sizer is based on characters.
44    /// let splitter = MarkdownSplitter::new(512);
45    /// ```
46    #[must_use]
47    pub fn new(chunk_config: impl Into<ChunkConfig<Sizer>>) -> Self {
48        Self {
49            chunk_config: chunk_config.into(),
50        }
51    }
52}
53
54impl<Sizer: ChunkSizer> MarkdownSplitter<Sizer> {
55    /// Generate a list of chunks from a given text. Each chunk will be up to
56    /// the `max_chunk_size`.
57    ///
58    /// ## Method
59    ///
60    /// To preserve as much semantic meaning within a chunk as possible, each chunk is composed of the largest semantic units that can fit in the next given chunk. For each splitter type, there is a defined set of semantic levels. Here is an example of the steps used:
61    ///
62    /// 1. Characters
63    /// 2. [Unicode Grapheme Cluster Boundaries](https://www.unicode.org/reports/tr29/#Grapheme_Cluster_Boundaries)
64    /// 3. [Unicode Word Boundaries](https://www.unicode.org/reports/tr29/#Word_Boundaries)
65    /// 4. [Unicode Sentence Boundaries](https://www.unicode.org/reports/tr29/#Sentence_Boundaries)
66    /// 5. Soft line breaks (single newline) which isn't necessarily a new element in Markdown.
67    /// 6. Inline elements such as: text nodes, emphasis, strong, strikethrough, link, image, table cells, inline code, footnote references, task list markers, and inline html.
68    /// 7. Block elements suce as: paragraphs, code blocks, footnote definitions, metadata. Also, a block quote or row/item within a table or list that can contain other "block" type elements, and a list or table that contains items.
69    /// 8. Thematic breaks or horizontal rules.
70    /// 9. Headings by level
71    ///
72    /// Splitting doesn't occur below the character level, otherwise you could get partial bytes of a char, which may not be a valid unicode str.
73    ///
74    /// Markdown is parsed according to the Commonmark spec, along with some optional features such as GitHub Flavored Markdown.
75    ///
76    /// ```
77    /// use text_splitter::MarkdownSplitter;
78    ///
79    /// let splitter = MarkdownSplitter::new(10);
80    /// let text = "# Header\n\nfrom a\ndocument";
81    /// let chunks = splitter.chunks(text).collect::<Vec<_>>();
82    ///
83    /// assert_eq!(vec!["# Header", "from a", "document"], chunks);
84    /// ```
85    pub fn chunks<'splitter, 'text: 'splitter>(
86        &'splitter self,
87        text: &'text str,
88    ) -> impl Iterator<Item = &'text str> + 'splitter {
89        Splitter::<_>::try_chunks(self, text).map(super::unwrap_infallible)
90    }
91
92    /// Returns an iterator over chunks of the text and their byte offsets.
93    /// Each chunk will be up to the `max_chunk_size`.
94    ///
95    /// See [`MarkdownSplitter::chunks`] for more information.
96    ///
97    /// ```
98    /// use text_splitter::MarkdownSplitter;
99    ///
100    /// let splitter = MarkdownSplitter::new(10);
101    /// let text = "# Header\n\nfrom a\ndocument";
102    /// let chunks = splitter.chunk_indices(text).collect::<Vec<_>>();
103    ///
104    /// assert_eq!(vec![(0, "# Header"), (10, "from a"), (17, "document")], chunks);
105    /// ```
106    pub fn chunk_indices<'splitter, 'text: 'splitter>(
107        &'splitter self,
108        text: &'text str,
109    ) -> impl Iterator<Item = (usize, &'text str)> + 'splitter {
110        Splitter::<_>::try_chunk_indices(self, text).map(super::unwrap_infallible)
111    }
112
113    /// Returns an iterator over chunks of the text with their byte and character offsets.
114    /// Each chunk will be up to the `chunk_capacity`.
115    ///
116    /// See [`MarkdownSplitter::chunks`] for more information.
117    ///
118    /// This will be more expensive than just byte offsets, and for most usage in Rust, just
119    /// having byte offsets is sufficient. But when interfacing with other languages or systems
120    /// that require character offsets, this will track the character offsets for you,
121    /// accounting for any trimming that may have occurred.
122    ///
123    /// ```
124    /// use text_splitter::{ChunkCharIndex, MarkdownSplitter};
125    ///
126    /// let splitter = MarkdownSplitter::new(10);
127    /// let text = "# Header\n\nfrom a\ndocument";
128    /// let chunks = splitter.chunk_char_indices(text).collect::<Vec<_>>();
129    ///
130    /// assert_eq!(vec![ChunkCharIndex { chunk: "# Header", byte_offset: 0, char_offset: 0 }, ChunkCharIndex { chunk: "from a", byte_offset: 10, char_offset: 10 }, ChunkCharIndex { chunk: "document", byte_offset: 17, char_offset: 17 }], chunks);
131    /// ```
132    pub fn chunk_char_indices<'splitter, 'text: 'splitter>(
133        &'splitter self,
134        text: &'text str,
135    ) -> impl Iterator<Item = ChunkCharIndex<'text>> + 'splitter {
136        Splitter::<_>::try_chunk_char_indices(self, text).map(super::unwrap_infallible)
137    }
138}
139
140impl<Sizer: TryChunkSizer> MarkdownSplitter<Sizer> {
141    /// Fallible version of [`Self::chunks`] for a [`TryChunkSizer`].
142    ///
143    /// Yields the first sizing error, then ends without calling the sizer again.
144    /// Collect into `Result<Vec<_>, _>` to stop at that error and discard partial
145    /// output. See [`TryChunkSizer`] for an example.
146    pub fn try_chunks<'splitter, 'text: 'splitter>(
147        &'splitter self,
148        text: &'text str,
149    ) -> impl Iterator<Item = Result<&'text str, Sizer::Error>> + 'splitter {
150        Splitter::<_>::try_chunks(self, text)
151    }
152
153    /// Fallible version of [`Self::chunk_indices`], including byte offsets.
154    ///
155    /// Yields the first sizing error, then ends without calling the sizer again.
156    pub fn try_chunk_indices<'splitter, 'text: 'splitter>(
157        &'splitter self,
158        text: &'text str,
159    ) -> impl Iterator<Item = Result<(usize, &'text str), Sizer::Error>> + 'splitter {
160        Splitter::<_>::try_chunk_indices(self, text)
161    }
162
163    /// Fallible version of [`Self::chunk_char_indices`], including byte and
164    /// character offsets.
165    ///
166    /// Yields the first sizing error, then ends without calling the sizer again.
167    pub fn try_chunk_char_indices<'splitter, 'text: 'splitter>(
168        &'splitter self,
169        text: &'text str,
170    ) -> impl Iterator<Item = Result<ChunkCharIndex<'text>, Sizer::Error>> + 'splitter {
171        Splitter::<_>::try_chunk_char_indices(self, text)
172    }
173}
174
175impl<Sizer> Splitter<Sizer> for MarkdownSplitter<Sizer>
176where
177    Sizer: TryChunkSizer,
178{
179    type Level = Element;
180
181    const TRIM: Trim = Trim::PreserveIndentation;
182
183    fn chunk_config(&self) -> &ChunkConfig<Sizer> {
184        &self.chunk_config
185    }
186
187    fn parse(&self, text: &str) -> Vec<(Self::Level, Range<usize>)> {
188        Parser::new_ext(text, Options::all())
189            .into_offset_iter()
190            .filter_map(|(event, range)| match event {
191                Event::Start(
192                    Tag::Emphasis
193                    | Tag::Strong
194                    | Tag::Strikethrough
195                    | Tag::Link { .. }
196                    | Tag::Image { .. }
197                    | Tag::Subscript
198                    | Tag::Superscript
199                    | Tag::TableCell,
200                )
201                | Event::Text(_)
202                | Event::HardBreak
203                | Event::Code(_)
204                | Event::InlineHtml(_)
205                | Event::InlineMath(_)
206                | Event::FootnoteReference(_)
207                | Event::TaskListMarker(_) => Some((Element::Inline, range)),
208                Event::SoftBreak => Some((Element::SoftBreak, range)),
209                Event::Html(_)
210                | Event::DisplayMath(_)
211                | Event::Start(
212                    Tag::Paragraph
213                    | Tag::CodeBlock(_)
214                    | Tag::FootnoteDefinition(_)
215                    | Tag::MetadataBlock(_)
216                    | Tag::TableHead
217                    | Tag::BlockQuote(_)
218                    | Tag::TableRow
219                    | Tag::Item
220                    | Tag::HtmlBlock
221                    | Tag::List(_)
222                    | Tag::Table(_)
223                    | Tag::DefinitionList
224                    | Tag::DefinitionListTitle
225                    | Tag::DefinitionListDefinition,
226                ) => Some((Element::Block, range)),
227                Event::Rule => Some((Element::Rule, range)),
228                Event::Start(Tag::Heading { level, .. }) => {
229                    Some((Element::Heading(level.into()), range))
230                }
231                // End events are identical to start, so no need to grab them.
232                Event::End(_) => None,
233            })
234            .collect()
235    }
236}
237
238/// Heading levels in markdown.
239/// Sorted in reverse order for sorting purposes.
240#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
241pub enum HeadingLevel {
242    H6,
243    H5,
244    H4,
245    H3,
246    H2,
247    H1,
248}
249
250impl From<pulldown_cmark::HeadingLevel> for HeadingLevel {
251    fn from(value: pulldown_cmark::HeadingLevel) -> Self {
252        match value {
253            pulldown_cmark::HeadingLevel::H1 => HeadingLevel::H1,
254            pulldown_cmark::HeadingLevel::H2 => HeadingLevel::H2,
255            pulldown_cmark::HeadingLevel::H3 => HeadingLevel::H3,
256            pulldown_cmark::HeadingLevel::H4 => HeadingLevel::H4,
257            pulldown_cmark::HeadingLevel::H5 => HeadingLevel::H5,
258            pulldown_cmark::HeadingLevel::H6 => HeadingLevel::H6,
259        }
260    }
261}
262
263/// How a particular semantic level relates to surrounding text elements.
264#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
265enum SemanticSplitPosition {
266    /// The semantic level should be treated as its own chunk.
267    Own,
268    /// The semantic level should be included in the next chunk.
269    Next,
270}
271
272/// Different semantic levels that text can be split by.
273/// Each level provides a method of splitting text into chunks of a given level
274/// as well as a fallback in case a given fallback is too large.
275#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
276pub enum Element {
277    /// Single line break, which isn't necessarily a new element in Markdown
278    SoftBreak,
279    /// An inline element that is within a larger element such as a paragraph, but
280    /// more specific than a sentence.
281    Inline,
282    /// Paragraph, code block, metadata, a row/item within a table or list, block quote, that can contain other "block" type elements, List or table that contains items
283    Block,
284    /// thematic break/horizontal rule
285    Rule,
286    /// Heading levels in markdown
287    Heading(HeadingLevel),
288}
289
290impl Element {
291    fn split_position(self) -> SemanticSplitPosition {
292        match self {
293            Self::SoftBreak | Self::Block | Self::Rule | Self::Inline => SemanticSplitPosition::Own,
294            // Attach it to the next text
295            Self::Heading(_) => SemanticSplitPosition::Next,
296        }
297    }
298
299    fn treat_whitespace_as_previous(self) -> bool {
300        match self {
301            Self::SoftBreak | Self::Inline | Self::Rule | Self::Heading(_) => false,
302            Self::Block => true,
303        }
304    }
305}
306
307impl SemanticLevel for Element {
308    fn sections(
309        text: &str,
310        level_ranges: impl Iterator<Item = (Self, Range<usize>)>,
311    ) -> impl Iterator<Item = (usize, &str)> {
312        let mut cursor = 0;
313        let mut final_match = false;
314        level_ranges
315            .batching(move |it| {
316                loop {
317                    match it.next() {
318                        // If we've hit the end, actually return None
319                        None if final_match => return None,
320                        // First time we hit None, return the final section of the text
321                        None => {
322                            final_match = true;
323                            return text.get(cursor..).map(|t| Either::Left(once((cursor, t))));
324                        }
325                        // Return text preceding match + the match
326                        Some((level, range)) => {
327                            let offset = cursor;
328                            match level.split_position() {
329                                SemanticSplitPosition::Own => {
330                                    if range.start < cursor {
331                                        continue;
332                                    }
333                                    let prev_section = text
334                                        .get(cursor..range.start)
335                                        .expect("invalid character sequence");
336                                    if level.treat_whitespace_as_previous()
337                                        && prev_section.chars().all(char::is_whitespace)
338                                    {
339                                        let section = text
340                                            .get(cursor..range.end)
341                                            .expect("invalid character sequence");
342                                        cursor = range.end;
343                                        return Some(Either::Left(once((offset, section))));
344                                    }
345                                    let separator = text
346                                        .get(range.start..range.end)
347                                        .expect("invalid character sequence");
348                                    cursor = range.end;
349                                    return Some(Either::Right(
350                                        [(offset, prev_section), (range.start, separator)]
351                                            .into_iter(),
352                                    ));
353                                }
354                                SemanticSplitPosition::Next => {
355                                    if range.start < cursor {
356                                        continue;
357                                    }
358                                    let prev_section = text
359                                        .get(cursor..range.start)
360                                        .expect("invalid character sequence");
361                                    // Separator will be part of the next chunk
362                                    cursor = range.start;
363                                    return Some(Either::Left(once((offset, prev_section))));
364                                }
365                            }
366                        }
367                    }
368                }
369            })
370            .flatten()
371            .filter(|(_, s)| !s.is_empty())
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use std::cmp::min;
378
379    use fake::{Fake, Faker};
380
381    use crate::splitter::SemanticSplitRanges;
382
383    use super::*;
384
385    #[test]
386    fn returns_one_chunk_if_text_is_shorter_than_max_chunk_size() {
387        let text = Faker.fake::<String>();
388        let chunks = MarkdownSplitter::new(ChunkConfig::new(text.chars().count()).with_trim(false))
389            .chunks(&text)
390            .collect::<Vec<_>>();
391
392        assert_eq!(vec![&text], chunks);
393    }
394
395    #[test]
396    fn returns_two_chunks_if_text_is_longer_than_max_chunk_size() {
397        let text1 = Faker.fake::<String>();
398        let text2 = Faker.fake::<String>();
399        let text = format!("{text1}{text2}");
400        // Round up to one above half so it goes to 2 chunks
401        let max_chunk_size = text.chars().count() / 2 + 1;
402
403        let chunks = MarkdownSplitter::new(ChunkConfig::new(max_chunk_size).with_trim(false))
404            .chunks(&text)
405            .collect::<Vec<_>>();
406
407        assert!(chunks.iter().all(|c| c.chars().count() <= max_chunk_size));
408
409        // Check that beginning of first chunk and text 1 matches
410        let len = min(text1.len(), chunks[0].len());
411        assert_eq!(text1[..len], chunks[0][..len]);
412        // Check that end of second chunk and text 2 matches
413        let len = min(text2.len(), chunks[1].len());
414        assert_eq!(
415            text2[(text2.len() - len)..],
416            chunks[1][chunks[1].len() - len..]
417        );
418
419        assert_eq!(chunks.join(""), text);
420    }
421
422    #[test]
423    fn empty_string() {
424        let text = "";
425        let chunks = MarkdownSplitter::new(ChunkConfig::new(100).with_trim(false))
426            .chunks(text)
427            .collect::<Vec<_>>();
428
429        assert!(chunks.is_empty());
430    }
431
432    #[test]
433    fn can_handle_unicode_characters() {
434        let text = "éé"; // Char that is more than one byte
435        let chunks = MarkdownSplitter::new(ChunkConfig::new(1).with_trim(false))
436            .chunks(text)
437            .collect::<Vec<_>>();
438
439        assert_eq!(vec!["é", "é"], chunks);
440    }
441
442    #[test]
443    fn chunk_by_graphemes() {
444        let text = "a̐éö̲\r\n";
445        let chunks = MarkdownSplitter::new(ChunkConfig::new(3).with_trim(false))
446            .chunks(text)
447            .collect::<Vec<_>>();
448
449        // \r\n is grouped together not separated
450        assert_eq!(vec!["a̐é", "ö̲", "\r\n"], chunks);
451    }
452
453    #[test]
454    fn trim_char_indices() {
455        let text = " a b ";
456        let chunks = MarkdownSplitter::new(1)
457            .chunk_indices(text)
458            .collect::<Vec<_>>();
459
460        assert_eq!(vec![(1, "a"), (3, "b")], chunks);
461    }
462
463    #[test]
464    fn chunk_char_indices() {
465        let text = " a b ";
466        let chunks = MarkdownSplitter::new(1)
467            .chunk_char_indices(text)
468            .collect::<Vec<_>>();
469
470        assert_eq!(
471            vec![
472                ChunkCharIndex {
473                    chunk: "a",
474                    byte_offset: 1,
475                    char_offset: 1
476                },
477                ChunkCharIndex {
478                    chunk: "b",
479                    byte_offset: 3,
480                    char_offset: 3,
481                },
482            ],
483            chunks
484        );
485    }
486
487    #[test]
488    fn graphemes_fallback_to_chars() {
489        let text = "a̐éö̲\r\n";
490        let chunks = MarkdownSplitter::new(ChunkConfig::new(1).with_trim(false))
491            .chunks(text)
492            .collect::<Vec<_>>();
493
494        assert_eq!(
495            vec!["a", "\u{310}", "é", "ö", "\u{332}", "\r", "\n"],
496            chunks
497        );
498    }
499
500    #[test]
501    fn trim_grapheme_indices() {
502        let text = "\r\na̐éö̲\r\n";
503        let chunks = MarkdownSplitter::new(3)
504            .chunk_indices(text)
505            .collect::<Vec<_>>();
506
507        assert_eq!(vec![(2, "a̐é"), (7, "ö̲")], chunks);
508    }
509
510    #[test]
511    fn grapheme_char_indices() {
512        let text = "\r\na̐éö̲\r\n";
513        let chunks = MarkdownSplitter::new(3)
514            .chunk_char_indices(text)
515            .collect::<Vec<_>>();
516
517        assert_eq!(
518            vec![
519                ChunkCharIndex {
520                    chunk: "a̐é",
521                    byte_offset: 2,
522                    char_offset: 2
523                },
524                ChunkCharIndex {
525                    chunk: "ö̲",
526                    byte_offset: 7,
527                    char_offset: 5
528                }
529            ],
530            chunks
531        );
532    }
533
534    #[test]
535    fn chunk_by_words() {
536        let text = "The quick brown fox can jump 32.3 feet, right?";
537        let chunks = MarkdownSplitter::new(ChunkConfig::new(10).with_trim(false))
538            .chunks(text)
539            .collect::<Vec<_>>();
540
541        assert_eq!(
542            vec![
543                "The quick ",
544                "brown fox ",
545                "can jump ",
546                "32.3 feet,",
547                " right?"
548            ],
549            chunks
550        );
551    }
552
553    #[test]
554    fn words_fallback_to_graphemes() {
555        let text = "Thé quick\r\n";
556        let chunks = MarkdownSplitter::new(ChunkConfig::new(2).with_trim(false))
557            .chunks(text)
558            .collect::<Vec<_>>();
559
560        assert_eq!(vec!["Th", "é ", "qu", "ic", "k", "\r\n"], chunks);
561    }
562
563    #[test]
564    fn trim_word_indices() {
565        let text = "Some text from a document";
566        let chunks = MarkdownSplitter::new(10)
567            .chunk_indices(text)
568            .collect::<Vec<_>>();
569
570        assert_eq!(
571            vec![(0, "Some text"), (10, "from a"), (17, "document")],
572            chunks
573        );
574    }
575
576    #[test]
577    fn chunk_by_sentences() {
578        let text = "Mr. Fox jumped. The dog was too lazy.";
579        let chunks = MarkdownSplitter::new(ChunkConfig::new(21).with_trim(false))
580            .chunks(text)
581            .collect::<Vec<_>>();
582
583        assert_eq!(vec!["Mr. Fox jumped. ", "The dog was too lazy."], chunks);
584    }
585
586    #[test]
587    fn sentences_falls_back_to_words() {
588        let text = "Mr. Fox jumped. The dog was too lazy.";
589        let chunks = MarkdownSplitter::new(ChunkConfig::new(16).with_trim(false))
590            .chunks(text)
591            .collect::<Vec<_>>();
592
593        assert_eq!(
594            vec!["Mr. Fox jumped. ", "The dog was too ", "lazy."],
595            chunks
596        );
597    }
598
599    #[test]
600    fn trim_sentence_indices() {
601        let text = "Some text. From a document.";
602        let chunks = MarkdownSplitter::new(10)
603            .chunk_indices(text)
604            .collect::<Vec<_>>();
605
606        assert_eq!(
607            vec![(0, "Some text."), (11, "From a"), (18, "document.")],
608            chunks
609        );
610    }
611
612    #[test]
613    fn test_no_markdown_separators() {
614        let splitter = MarkdownSplitter::new(10);
615        let markdown =
616            SemanticSplitRanges::new(splitter.parse("Some text without any markdown separators"));
617
618        assert_eq!(
619            vec![(Element::Block, 0..41), (Element::Inline, 0..41)],
620            markdown.ranges_after_offset(0).collect::<Vec<_>>()
621        );
622    }
623
624    #[test]
625    fn test_checklist() {
626        let splitter = MarkdownSplitter::new(10);
627        let markdown =
628            SemanticSplitRanges::new(splitter.parse("- [ ] incomplete task\n- [x] completed task"));
629
630        assert_eq!(
631            vec![
632                (Element::Block, 0..42),
633                (Element::Block, 0..22),
634                (Element::Inline, 2..5),
635                (Element::Inline, 6..21),
636                (Element::Block, 22..42),
637                (Element::Inline, 24..27),
638                (Element::Inline, 28..42),
639            ],
640            markdown.ranges_after_offset(0).collect::<Vec<_>>()
641        );
642    }
643
644    #[test]
645    fn test_footnote_reference() {
646        let splitter = MarkdownSplitter::new(10);
647        let markdown = SemanticSplitRanges::new(splitter.parse("Footnote[^1]"));
648
649        assert_eq!(
650            vec![
651                (Element::Block, 0..12),
652                (Element::Inline, 0..8),
653                (Element::Inline, 8..12),
654            ],
655            markdown.ranges_after_offset(0).collect::<Vec<_>>()
656        );
657    }
658
659    #[test]
660    fn test_inline_code() {
661        let splitter = MarkdownSplitter::new(10);
662        let markdown = SemanticSplitRanges::new(splitter.parse("`bash`"));
663
664        assert_eq!(
665            vec![(Element::Block, 0..6), (Element::Inline, 0..6)],
666            markdown.ranges_after_offset(0).collect::<Vec<_>>()
667        );
668    }
669
670    #[test]
671    fn test_emphasis() {
672        let splitter = MarkdownSplitter::new(10);
673        let markdown = SemanticSplitRanges::new(splitter.parse("*emphasis*"));
674
675        assert_eq!(
676            vec![
677                (Element::Block, 0..10),
678                (Element::Inline, 0..10),
679                (Element::Inline, 1..9),
680            ],
681            markdown.ranges_after_offset(0).collect::<Vec<_>>()
682        );
683    }
684
685    #[test]
686    fn test_strong() {
687        let splitter = MarkdownSplitter::new(10);
688        let markdown = SemanticSplitRanges::new(splitter.parse("**emphasis**"));
689
690        assert_eq!(
691            vec![
692                (Element::Block, 0..12),
693                (Element::Inline, 0..12),
694                (Element::Inline, 2..10),
695            ],
696            markdown.ranges_after_offset(0).collect::<Vec<_>>()
697        );
698    }
699
700    #[test]
701    fn test_strikethrough() {
702        let splitter = MarkdownSplitter::new(10);
703        let markdown = SemanticSplitRanges::new(splitter.parse("~~emphasis~~"));
704
705        assert_eq!(
706            vec![
707                (Element::Block, 0..12),
708                (Element::Inline, 0..12),
709                (Element::Inline, 2..10),
710            ],
711            markdown.ranges_after_offset(0).collect::<Vec<_>>()
712        );
713    }
714
715    #[test]
716    fn test_link() {
717        let splitter = MarkdownSplitter::new(10);
718        let markdown = SemanticSplitRanges::new(splitter.parse("[link](url)"));
719
720        assert_eq!(
721            vec![
722                (Element::Block, 0..11),
723                (Element::Inline, 0..11),
724                (Element::Inline, 1..5),
725            ],
726            markdown.ranges_after_offset(0).collect::<Vec<_>>()
727        );
728    }
729
730    #[test]
731    fn test_image() {
732        let splitter = MarkdownSplitter::new(10);
733        let markdown = SemanticSplitRanges::new(splitter.parse("![link](url)"));
734
735        assert_eq!(
736            vec![
737                (Element::Block, 0..12),
738                (Element::Inline, 0..12),
739                (Element::Inline, 2..6),
740            ],
741            markdown.ranges_after_offset(0).collect::<Vec<_>>()
742        );
743    }
744
745    #[test]
746    fn test_inline_html() {
747        let splitter = MarkdownSplitter::new(10);
748        let markdown = SemanticSplitRanges::new(splitter.parse("<span>Some text</span>"));
749
750        assert_eq!(
751            vec![
752                (Element::Block, 0..22),
753                (Element::Inline, 0..6),
754                (Element::Inline, 6..15),
755                (Element::Inline, 15..22),
756            ],
757            markdown.ranges_after_offset(0).collect::<Vec<_>>()
758        );
759    }
760
761    #[test]
762    fn test_html() {
763        let splitter = MarkdownSplitter::new(10);
764        let markdown = SemanticSplitRanges::new(splitter.parse("<div>Some text</div>"));
765
766        assert_eq!(
767            vec![(Element::Block, 0..20), (Element::Block, 0..20)],
768            markdown.ranges_after_offset(0).collect::<Vec<_>>()
769        );
770    }
771
772    #[test]
773    fn test_table() {
774        let splitter = MarkdownSplitter::new(10);
775        let markdown = SemanticSplitRanges::new(
776            splitter.parse("| Header 1 | Header 2 |\n| --- | --- |\n| Cell 1 | Cell 2 |"),
777        );
778        assert_eq!(
779            vec![
780                (Element::Block, 0..57),
781                (Element::Block, 0..24),
782                (Element::Inline, 1..11),
783                (Element::Inline, 2..10),
784                (Element::Inline, 12..22),
785                (Element::Inline, 13..21),
786                (Element::Block, 38..57),
787                (Element::Inline, 39..47),
788                (Element::Inline, 40..46),
789                (Element::Inline, 48..56),
790                (Element::Inline, 49..55)
791            ],
792            markdown.ranges_after_offset(0).collect::<Vec<_>>()
793        );
794    }
795
796    #[test]
797    fn test_softbreak() {
798        let splitter = MarkdownSplitter::new(10);
799        let markdown = SemanticSplitRanges::new(splitter.parse("Some text\nwith a softbreak"));
800
801        assert_eq!(
802            vec![
803                (Element::Block, 0..26),
804                (Element::Inline, 0..9),
805                (Element::SoftBreak, 9..10),
806                (Element::Inline, 10..26)
807            ],
808            markdown.ranges_after_offset(0).collect::<Vec<_>>()
809        );
810    }
811
812    #[test]
813    fn test_hardbreak() {
814        let splitter = MarkdownSplitter::new(10);
815        let markdown = SemanticSplitRanges::new(splitter.parse("Some text\\\nwith a hardbreak"));
816
817        assert_eq!(
818            vec![
819                (Element::Block, 0..27),
820                (Element::Inline, 0..9),
821                (Element::Inline, 9..11),
822                (Element::Inline, 11..27)
823            ],
824            markdown.ranges_after_offset(0).collect::<Vec<_>>()
825        );
826    }
827
828    #[test]
829    fn test_footnote_def() {
830        let splitter = MarkdownSplitter::new(10);
831        let markdown = SemanticSplitRanges::new(splitter.parse("[^first]: Footnote"));
832
833        assert_eq!(
834            vec![
835                (Element::Block, 0..18),
836                (Element::Block, 10..18),
837                (Element::Inline, 10..18)
838            ],
839            markdown.ranges_after_offset(0).collect::<Vec<_>>()
840        );
841    }
842
843    #[test]
844    fn test_code_block() {
845        let splitter = MarkdownSplitter::new(10);
846        let markdown = SemanticSplitRanges::new(splitter.parse("```\ncode\n```"));
847
848        assert_eq!(
849            vec![(Element::Block, 0..12), (Element::Inline, 4..9)],
850            markdown.ranges_after_offset(0).collect::<Vec<_>>()
851        );
852    }
853
854    #[test]
855    fn test_block_quote() {
856        let splitter = MarkdownSplitter::new(10);
857        let markdown = SemanticSplitRanges::new(splitter.parse("> quote"));
858
859        assert_eq!(
860            vec![
861                (Element::Block, 0..7),
862                (Element::Block, 2..7),
863                (Element::Inline, 2..7)
864            ],
865            markdown.ranges_after_offset(0).collect::<Vec<_>>()
866        );
867    }
868
869    #[test]
870    fn test_with_rule() {
871        let splitter = MarkdownSplitter::new(10);
872        let markdown = SemanticSplitRanges::new(splitter.parse("Some text\n\n---\n\nwith a rule"));
873
874        assert_eq!(
875            vec![
876                (Element::Block, 0..10),
877                (Element::Inline, 0..9),
878                (Element::Rule, 11..15),
879                (Element::Block, 16..27),
880                (Element::Inline, 16..27)
881            ],
882            markdown.ranges_after_offset(0).collect::<Vec<_>>()
883        );
884    }
885
886    #[test]
887    fn test_heading() {
888        for (index, (heading, level)) in [
889            ("#", HeadingLevel::H1),
890            ("##", HeadingLevel::H2),
891            ("###", HeadingLevel::H3),
892            ("####", HeadingLevel::H4),
893            ("#####", HeadingLevel::H5),
894            ("######", HeadingLevel::H6),
895        ]
896        .into_iter()
897        .enumerate()
898        {
899            let splitter = MarkdownSplitter::new(10);
900            let markdown = SemanticSplitRanges::new(splitter.parse(&format!("{heading} Heading")));
901
902            assert_eq!(
903                vec![
904                    (Element::Heading(level), 0..9 + index),
905                    (Element::Inline, 2 + index..9 + index)
906                ],
907                markdown.ranges_after_offset(0).collect::<Vec<_>>()
908            );
909        }
910    }
911
912    #[test]
913    fn test_ranges_after_offset_block() {
914        let splitter = MarkdownSplitter::new(10);
915        let markdown =
916            SemanticSplitRanges::new(splitter.parse("- [ ] incomplete task\n- [x] completed task"));
917
918        assert_eq!(
919            vec![(Element::Block, 0..22), (Element::Block, 22..42),],
920            markdown
921                .level_ranges_after_offset(0, Element::Block)
922                .collect::<Vec<_>>()
923        );
924    }
925}
text_splitter/splitter/markdown.rs

text_splitter/splitter/
markdown.rs
