oxidize_pdf/page_labels/
page_label_tree.rs

1//! Page label tree structure for managing page numbering
2
3use crate::objects::{Array, Dictionary, Object};
4use crate::page_labels::PageLabel;
5use std::collections::BTreeMap;
6
7/// Page label tree - manages custom page numbering for a document
8#[derive(Debug, Clone)]
9pub struct PageLabelTree {
10    /// Page label ranges, sorted by starting page
11    ranges: BTreeMap<u32, PageLabel>,
12}
13
14impl PageLabelTree {
15    /// Create a new empty page label tree
16    pub fn new() -> Self {
17        Self {
18            ranges: BTreeMap::new(),
19        }
20    }
21
22    /// Add a page label range
23    pub fn add_range(&mut self, start_page: u32, label: PageLabel) {
24        self.ranges.insert(start_page, label);
25    }
26
27    /// Get the page label for a specific page
28    pub fn get_label(&self, page_index: u32) -> Option<String> {
29        // Find the applicable range
30        let mut applicable_range = None;
31        let mut range_start = 0;
32
33        for (&start, label) in &self.ranges {
34            if start <= page_index {
35                applicable_range = Some(label);
36                range_start = start;
37            } else {
38                break;
39            }
40        }
41
42        // Format the label if found
43        applicable_range.map(|label| {
44            let offset = page_index - range_start;
45            label.format_label(offset)
46        })
47    }
48
49    /// Get all page labels for a document
50    pub fn get_all_labels(&self, total_pages: u32) -> Vec<String> {
51        (0..total_pages)
52            .map(|i| self.get_label(i).unwrap_or_else(|| (i + 1).to_string()))
53            .collect()
54    }
55
56    /// Convert to PDF number tree dictionary
57    pub fn to_dict(&self) -> Dictionary {
58        let mut dict = Dictionary::new();
59
60        // Create nums array [key1 val1 key2 val2 ...]
61        let mut nums = Array::new();
62
63        for (&start_page, label) in &self.ranges {
64            nums.push(Object::Integer(start_page as i64));
65            nums.push(Object::Dictionary(label.to_dict()));
66        }
67
68        dict.set("Nums", Object::Array(nums.into()));
69
70        dict
71    }
72
73    /// Create from PDF dictionary
74    pub fn from_dict(dict: &Dictionary) -> Option<Self> {
75        let nums_array = match dict.get("Nums")? {
76            Object::Array(arr) => arr,
77            _ => return None,
78        };
79        let mut tree = Self::new();
80
81        // Parse pairs of [page_index, label_dict]
82        let elements: Vec<&Object> = nums_array.iter().collect();
83        for i in (0..elements.len()).step_by(2) {
84            if i + 1 >= elements.len() {
85                break;
86            }
87
88            let page_index = match elements[i] {
89                Object::Integer(n) => *n as u32,
90                _ => continue,
91            };
92            let label_dict = match elements[i + 1] {
93                Object::Dictionary(d) => d,
94                _ => continue,
95            };
96
97            // Parse label from dictionary
98            let style = if let Some(Object::Name(type_name)) = label_dict.get("Type") {
99                match type_name.as_str() {
100                    "D" => PageLabelStyle::DecimalArabic,
101                    "r" => PageLabelStyle::UppercaseRoman,
102                    "R" => PageLabelStyle::LowercaseRoman,
103                    "A" => PageLabelStyle::UppercaseLetters,
104                    "a" => PageLabelStyle::LowercaseLetters,
105                    _ => PageLabelStyle::None,
106                }
107            } else {
108                PageLabelStyle::None
109            };
110
111            let mut label = PageLabel::new(style);
112
113            if let Some(Object::String(prefix)) = label_dict.get("P") {
114                label = label.with_prefix(prefix);
115            }
116
117            if let Some(Object::Integer(start)) = label_dict.get("St") {
118                label = label.starting_at(*start as u32);
119            }
120
121            tree.add_range(page_index, label);
122        }
123
124        Some(tree)
125    }
126}
127
128impl Default for PageLabelTree {
129    fn default() -> Self {
130        Self::new()
131    }
132}
133
134/// Builder for creating page label trees
135pub struct PageLabelBuilder {
136    tree: PageLabelTree,
137    current_page: u32,
138}
139
140impl Default for PageLabelBuilder {
141    fn default() -> Self {
142        Self::new()
143    }
144}
145
146impl PageLabelBuilder {
147    /// Create a new page label builder
148    pub fn new() -> Self {
149        Self {
150            tree: PageLabelTree::new(),
151            current_page: 0,
152        }
153    }
154
155    /// Add a range with a specific label
156    pub fn add_range(mut self, num_pages: u32, label: PageLabel) -> Self {
157        self.tree.add_range(self.current_page, label);
158        self.current_page += num_pages;
159        self
160    }
161
162    /// Add pages with decimal numbering
163    pub fn decimal_pages(self, num_pages: u32) -> Self {
164        self.add_range(num_pages, PageLabel::decimal())
165    }
166
167    /// Add pages with roman numbering
168    pub fn roman_pages(self, num_pages: u32, uppercase: bool) -> Self {
169        let label = if uppercase {
170            PageLabel::roman_uppercase()
171        } else {
172            PageLabel::roman_lowercase()
173        };
174        self.add_range(num_pages, label)
175    }
176
177    /// Add pages with letter numbering
178    pub fn letter_pages(self, num_pages: u32, uppercase: bool) -> Self {
179        let label = if uppercase {
180            PageLabel::letters_uppercase()
181        } else {
182            PageLabel::letters_lowercase()
183        };
184        self.add_range(num_pages, label)
185    }
186
187    /// Add pages with only a prefix
188    pub fn prefix_pages(self, num_pages: u32, prefix: impl Into<String>) -> Self {
189        self.add_range(num_pages, PageLabel::prefix_only(prefix))
190    }
191
192    /// Build the page label tree
193    pub fn build(self) -> PageLabelTree {
194        self.tree
195    }
196}
197
198// Import PageLabelStyle from the other module
199use crate::page_labels::PageLabelStyle;
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn test_page_label_tree() {
207        let mut tree = PageLabelTree::new();
208
209        // Add roman numerals for first 3 pages
210        tree.add_range(0, PageLabel::roman_lowercase());
211
212        // Add decimal starting at page 3
213        tree.add_range(3, PageLabel::decimal());
214
215        // Test labels
216        assert_eq!(tree.get_label(0), Some("i".to_string()));
217        assert_eq!(tree.get_label(1), Some("ii".to_string()));
218        assert_eq!(tree.get_label(2), Some("iii".to_string()));
219        assert_eq!(tree.get_label(3), Some("1".to_string()));
220        assert_eq!(tree.get_label(4), Some("2".to_string()));
221        assert_eq!(tree.get_label(5), Some("3".to_string()));
222    }
223
224    #[test]
225    fn test_page_label_with_prefix() {
226        let mut tree = PageLabelTree::new();
227
228        // Preface with prefix
229        tree.add_range(0, PageLabel::prefix_only("Cover"));
230        tree.add_range(1, PageLabel::roman_lowercase().with_prefix("p. "));
231        tree.add_range(4, PageLabel::decimal().with_prefix("Chapter "));
232
233        assert_eq!(tree.get_label(0), Some("Cover".to_string()));
234        assert_eq!(tree.get_label(1), Some("p. i".to_string()));
235        assert_eq!(tree.get_label(2), Some("p. ii".to_string()));
236        assert_eq!(tree.get_label(3), Some("p. iii".to_string()));
237        assert_eq!(tree.get_label(4), Some("Chapter 1".to_string()));
238        assert_eq!(tree.get_label(5), Some("Chapter 2".to_string()));
239    }
240
241    #[test]
242    fn test_page_label_with_start() {
243        let mut tree = PageLabelTree::new();
244
245        // Start numbering at 10
246        tree.add_range(0, PageLabel::decimal().starting_at(10));
247
248        assert_eq!(tree.get_label(0), Some("10".to_string()));
249        assert_eq!(tree.get_label(1), Some("11".to_string()));
250        assert_eq!(tree.get_label(2), Some("12".to_string()));
251    }
252
253    #[test]
254    fn test_get_all_labels() {
255        let mut tree = PageLabelTree::new();
256        tree.add_range(0, PageLabel::roman_lowercase());
257        tree.add_range(2, PageLabel::decimal());
258
259        let labels = tree.get_all_labels(5);
260        assert_eq!(labels, vec!["i", "ii", "1", "2", "3"]);
261    }
262
263    #[test]
264    fn test_page_label_builder() {
265        let tree = PageLabelBuilder::new()
266            .prefix_pages(1, "Cover")
267            .roman_pages(3, false)
268            .decimal_pages(10)
269            .letter_pages(3, true)
270            .build();
271
272        assert_eq!(tree.get_label(0), Some("Cover".to_string()));
273        assert_eq!(tree.get_label(1), Some("i".to_string()));
274        assert_eq!(tree.get_label(2), Some("ii".to_string()));
275        assert_eq!(tree.get_label(3), Some("iii".to_string()));
276        assert_eq!(tree.get_label(4), Some("1".to_string()));
277        assert_eq!(tree.get_label(13), Some("10".to_string()));
278        assert_eq!(tree.get_label(14), Some("A".to_string()));
279        assert_eq!(tree.get_label(15), Some("B".to_string()));
280        assert_eq!(tree.get_label(16), Some("C".to_string()));
281    }
282
283    #[test]
284    fn test_to_dict() {
285        let mut tree = PageLabelTree::new();
286        tree.add_range(0, PageLabel::roman_lowercase());
287        tree.add_range(3, PageLabel::decimal().with_prefix("Page "));
288
289        let dict = tree.to_dict();
290        assert!(dict.get("Nums").is_some());
291    }
292}
oxidize_pdf/page_labels/page_label_tree.rs

oxidize_pdf/page_labels/
page_label_tree.rs
