Article 1
+Some text with bold and emphasis
+
+ Hello & Welcome
+ +
+ `, ``, `
Hello
`) based on the configuration
+fn add_attributes(
+ config: &HtmlTransformerConfig,
+ element: &mut BytesStart,
+ is_root: bool,
+ captured_attributes: &mut Vec<(String, Vec)>,
+) {
+ let mut added_attrs = Vec::new();
+
+ // Add root attributes if this is a root element
+ if is_root {
+ for attr in &config.root_attributes {
+ element.push_attribute((attr.as_str(), ""));
+ added_attrs.push(attr.clone());
+ }
+ }
+
+ // Add attributes that should be applied to all elements
+ for attr in &config.all_attributes {
+ element.push_attribute((attr.as_str(), ""));
+ added_attrs.push(attr.clone());
+ }
+
+ // If we're watching for a specific attribute, check if this element has it
+ if let Some(watch_attr) = &config.watch_on_attribute {
+ if let Some(attr_value) = element
+ .attributes()
+ .find(|a| {
+ if let Ok(attr) = a {
+ String::from_utf8_lossy(attr.key.as_ref()) == *watch_attr
+ } else {
+ false
+ }
+ })
+ .and_then(|a| a.ok())
+ .map(|a| String::from_utf8_lossy(a.value.as_ref()).into_owned())
+ {
+ captured_attributes.push((attr_value, added_attrs));
+ }
+ }
+}
+
+/// Main entrypoint. Transform HTML by adding attributes to the elements.
+pub fn transform(
+ config: &HtmlTransformerConfig,
+ html: &str,
+) -> Result<(String, Vec<(String, Vec)>), Box> {
+ let mut reader = Reader::from_str(html);
+ let reader_config = reader.config_mut();
+ reader_config.check_end_names = config.check_end_names;
+
+ // We transform the HTML by reading it and writing it simultaneously
+ let mut writer = Writer::new(Cursor::new(Vec::new()));
+ let mut captured_attributes = Vec::new();
+
+ // Track the nesting depth of elements to identify root elements (depth == 0)
+ let mut depth: i32 = 0;
+
+ // Read the HTML event by event
+ loop {
+ match reader.read_event() {
+ // Start tag
+ Ok(Event::Start(e)) => {
+ let tag_name = String::from_utf8_lossy(e.name().as_ref())
+ .to_string()
+ .to_lowercase();
+ let mut elem = e.into_owned();
+ add_attributes(config, &mut elem, depth == 0, &mut captured_attributes);
+
+ // For void elements, write as Empty event
+ if config.void_elements.contains(&tag_name) {
+ writer.write_event(Event::Empty(elem))?;
+ } else {
+ writer.write_event(Event::Start(elem))?;
+ depth += 1;
+ }
+ }
+
+ // End tag
+ Ok(Event::End(e)) => {
+ let tag_name = String::from_utf8_lossy(e.name().as_ref())
+ .to_string()
+ .to_lowercase();
+
+ // Skip end tags for void elements
+ if !config.void_elements.contains(&tag_name) {
+ writer.write_event(Event::End(e))?;
+ depth -= 1;
+ }
+ }
+
+ // Empty element (AKA void or self-closing tag, e.g. `
`) + Ok(Event::Empty(e)) => { + let mut elem = e.into_owned(); + add_attributes(config, &mut elem, depth == 0, &mut captured_attributes); + writer.write_event(Event::Empty(elem))?; + } + + // End of file + Ok(Event::Eof) => break, + // Other events (e.g. comments, processing instructions, etc.) + Ok(e) => writer.write_event(e)?, + Err(e) => return Err(Box::new(e)), + } + } + + // Convert the transformed HTML to a string + let result = String::from_utf8(writer.into_inner().into_inner())?; + Ok((result, captured_attributes)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic_transformation() { + let config = HtmlTransformerConfig::new( + vec!["data-root".to_string()], + vec!["data-all".to_string()], + false, + None, + ); + + let input = "";
+ let (result, _) = transform(&config, input).unwrap();
+
+ assert!(result.contains("data-root"));
+ assert!(result.contains("data-all"));
+ }
+
+ #[test]
+ fn test_multiple_roots() {
+ let config = HtmlTransformerConfig::new(
+ vec!["data-root".to_string()],
+ vec!["data-all".to_string()],
+ false,
+ None,
+ );
+
+ let input = "Test Page
+
+
+ "#;
+
+ let (result, _) = transform(&config, input).unwrap();
+
+ // Check that it parsed successfully
+ assert!(result.contains(r#"Test Page"#));
+ assert!(result.contains(r#""));
+ assert!(!result.contains(""));
+ assert!(result.contains("/>"));
+ }
+
+ #[test]
+ fn test_config_check_end_names() {
+ // Test with check_end_names = false (lenient mode)
+ let config = HtmlTransformerConfig::new(
+ vec!["data-root".to_string()],
+ vec!["data-v-123".to_string()],
+ false, // Don't check end names
+ None,
+ );
+
+ // These should parse successfully with check_end_names = false
+ let lenient_cases = [
+ "", // Mismatched nesting
+ "
`) + Ok(Event::Empty(e)) => { + let mut elem = e.into_owned(); + add_attributes(config, &mut elem, depth == 0, &mut captured_attributes); + writer.write_event(Event::Empty(elem))?; + } + + // End of file + Ok(Event::Eof) => break, + // Other events (e.g. comments, processing instructions, etc.) + Ok(e) => writer.write_event(e)?, + Err(e) => return Err(Box::new(e)), + } + } + + // Convert the transformed HTML to a string + let result = String::from_utf8(writer.into_inner().into_inner())?; + Ok((result, captured_attributes)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic_transformation() { + let config = HtmlTransformerConfig::new( + vec!["data-root".to_string()], + vec!["data-all".to_string()], + false, + None, + ); + + let input = "
Hello
First
Second";
+ let (result, _) = transform(&config, input).unwrap();
+
+ // Both root elements should have data-root
+ assert_eq!(result.matches("data-root").count(), 2);
+ // All elements should have data-all
+ assert_eq!(result.matches("data-all").count(), 2);
+ }
+
+ #[test]
+ fn test_complex_html() {
+ let config = HtmlTransformerConfig::new(
+ vec!["data-root".to_string()],
+ vec!["data-all".to_string(), "data-v-123".to_string()],
+ false,
+ None,
+ );
+
+ let input = r#"
+
+
+
+
+
+
+
+
+
+
+ "#;
+
+ let (result, _) = transform(&config, input).unwrap();
+
+ // Check root elements have root attributes
+ assert!(result.contains(
+ r#"
+ "#
+ ));
+ assert!(result.contains(r#"
"));
+ }
+
+ #[test]
+ fn test_html_head_with_meta() {
+ let config = HtmlTransformerConfig::new(
+ vec!["data-root".to_string()],
+ vec!["data-v-123".to_string()],
+ false,
+ None,
+ );
+
+ let input = r#"
+
+
+ Hello
Text", // Wrong closing tag
+ "";
+ assert!(transform(&config, valid_input).is_ok());
+ }
+
+ #[test]
+ fn test_watch_attribute() {
+ let config = HtmlTransformerConfig::new(
+ vec!["data-root".to_string()],
+ vec!["data-v-123".to_string()],
+ false,
+ Some("data-id".to_string()),
+ );
+
+ let input = r#"
+
Text", // Non-matching end tag + ]; + + for input in lenient_cases { + assert!(transform(&config, input).is_ok()); + } + + // Test with check_end_names = true (strict mode) + let config = HtmlTransformerConfig::new( + vec!["data-root".to_string()], + vec!["data-v-123".to_string()], + true, // Check end names + None, + ); + + // These should fail with check_end_names = true + for input in lenient_cases { + assert!(transform(&config, input).is_err()); + } + + // But well-formed HTML should still work + let valid_input = "
Hello
+
+
"#;
+
+ let (result, captured) = transform(&config, input).unwrap();
+
+ println!("result: {}", result);
+ println!("captured: {:?}", captured);
+
+ // Verify HTML transformation worked
+ assert!(result.contains(r#"Regular element
+ Nested element +
+ "#));
+ assert!(result.contains(r#""#));
+ assert!(result.contains(r#""#));
+
+ // Verify attribute capturing
+ assert_eq!(captured.len(), 3);
+ assert!(captured.iter().any(|(id, attrs)| id == "123"
+ && attrs.contains(&"data-root".to_string())
+ && attrs.contains(&"data-v-123".to_string())));
+ assert!(captured
+ .iter()
+ .any(|(id, attrs)| id == "456" && attrs.contains(&"data-v-123".to_string())));
+ assert!(captured
+ .iter()
+ .any(|(id, attrs)| id == "789" && attrs.contains(&"data-v-123".to_string())));
+ }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..050fa3e
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,11 @@
+use pyo3::prelude::*;
+use pyo3::types::PyModule;
+
+mod html_parser;
+
+/// A Python module implemented in Rust for high-performance HTML transformation.
+#[pymodule]
+fn djc_core_html_parser(_py: Python, m: &PyModule) -> PyResult<()> {
+ m.add_function(wrap_pyfunction!(html_parser::set_html_attributes, m)?)?;
+ Ok(())
+}
diff --git a/tests/benchmark.py b/tests/benchmark.py
new file mode 100644
index 0000000..1cf02b1
--- /dev/null
+++ b/tests/benchmark.py
@@ -0,0 +1,97 @@
+from statistics import mean, stdev
+import time
+
+from djc_core_html_parser import set_html_attributes
+
+
+def generate_large_html(num_elements: int = 1000) -> str:
+ """Generate a large HTML document with various features for benchmarking."""
+ elements = []
+ for i in range(num_elements):
+ # Mix of different elements and features
+ if i % 5 == 0:
+ # Void element with multiple attributes
+ elements.append(f'
')
+ elif i % 5 == 1:
+ # Nested divs with attributes
+ elements.append(
+ f"""
+
+
+ """
+ )
+
+ return f"""
+
+
+
+ Benchmark Page
+
+
+
+ {''.join(elements)}
+
+
+ """
+
+
+# Generate test HTML
+HTML_SIZE = 27_000 # Set to 11_000 for 2MB
+NUM_ITER = 2
+html = generate_large_html(HTML_SIZE)
+print(f"\nBenchmarking with HTML size: {len(html) // 1_000} KB")
+
+root_attributes = ["data-root-id"]
+all_attributes = ["data-v-123"]
+
+# Test transform
+modify_times = []
+for i in range(NUM_ITER): # Run N iterations
+
+ start = time.perf_counter()
+ set_html_attributes(html, root_attributes, all_attributes, watch_on_attribute="data-id")
+ modify_time = time.perf_counter() - start
+ modify_times.append(modify_time)
+
+print("\nTransform:")
+print(f" Total: {sum(modify_times):.3f}s")
+print(f" Min: {min(modify_times):.3f}s")
+print(f" Max: {max(modify_times):.3f}s")
+print(f" Avg: {mean(modify_times):.3f}s")
+print(f" Std: {stdev(modify_times):.3f}s")
diff --git a/tests/test_html_parser.py b/tests/test_html_parser.py
new file mode 100644
index 0000000..57cf3b1
--- /dev/null
+++ b/tests/test_html_parser.py
@@ -0,0 +1,148 @@
+# This same set of tests is also found in django-components, to ensure that
+# this implementation can be replaced with the django-components' pure-python implementation
+
+from djc_core_html_parser import set_html_attributes
+from typing import Dict, List
+
+
+def test_basic_transformation():
+ html = ""
+ result, _ = set_html_attributes(html, ["data-root"], ["data-all"])
+ expected = ''
+ assert result == expected
+
+
+def test_multiple_roots():
+ html = "
", '
'),
+ ('', ''),
+ ]
+
+ for input_html, expected in test_cases:
+ result, _ = set_html_attributes(input_html, ["data-root"], ["data-v-123"])
+ assert result == expected
+
+
+def test_html_head_with_meta():
+ html = """
+
+
+ Test Page
+
+
+ """
+
+ result, _ = set_html_attributes(html, ["data-root"], ["data-v-123"])
+ expected = """
+
+
+ Test Page
+
+
+ """
+ assert result == expected
+
+
+def test_watch_attribute():
+ html = """
+
"#));
+
+ // Verify attribute capturing
+ assert_eq!(captured.len(), 3);
+ assert!(captured.iter().any(|(id, attrs)| id == "123"
+ && attrs.contains(&"data-root".to_string())
+ && attrs.contains(&"data-v-123".to_string())));
+ assert!(captured
+ .iter()
+ .any(|(id, attrs)| id == "456" && attrs.contains(&"data-v-123".to_string())));
+ assert!(captured
+ .iter()
+ .any(|(id, attrs)| id == "789" && attrs.contains(&"data-v-123".to_string())));
+ }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..050fa3e
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,11 @@
+use pyo3::prelude::*;
+use pyo3::types::PyModule;
+
+mod html_parser;
+
+/// A Python module implemented in Rust for high-performance HTML transformation.
+#[pymodule]
+fn djc_core_html_parser(_py: Python, m: &PyModule) -> PyResult<()> {
+ m.add_function(wrap_pyfunction!(html_parser::set_html_attributes, m)?)?;
+ Ok(())
+}
diff --git a/tests/benchmark.py b/tests/benchmark.py
new file mode 100644
index 0000000..1cf02b1
--- /dev/null
+++ b/tests/benchmark.py
@@ -0,0 +1,97 @@
+from statistics import mean, stdev
+import time
+
+from djc_core_html_parser import set_html_attributes
+
+
+def generate_large_html(num_elements: int = 1000) -> str:
+ """Generate a large HTML document with various features for benchmarking."""
+ elements = []
+ for i in range(num_elements):
+ # Mix of different elements and features
+ if i % 5 == 0:
+ # Void element with multiple attributes
+ elements.append(f'')
+ elif i % 5 == 1:
+ # Nested divs with attributes
+ elements.append(
+ f"""
+
+
+ """
+ )
+ elif i % 5 == 2:
+ # Script tag with content
+ elements.append(
+ f"""
+
+ """
+ )
+ elif i % 5 == 3:
+ # CDATA section
+ elements.append(
+ f"""
+
+ ]]>
+ """
+ )
+ else:
+ # Regular element with attributes
+ elements.append(
+ f"""
+
+
+ Content {i}
+ +Heading {i}
+Paragraph {i}
+Hello
Hello
First
Second"
+ result, _ = set_html_attributes(html, ["data-root"], ["data-all"])
+ expected = 'First
Second'
+ assert result == expected
+
+
+def test_complex_html():
+ html = """
+
+
+
+
+
+
+
+
+
+
+ """
+
+ result, _ = set_html_attributes(html, ["data-root"], ["data-all", "data-v-123"])
+ expected = """
+ Hello & Welcome
+ +Article 1
+Some text with bold and emphasis
+
+
+
+
+
+
+
+
+
+
+
+ """ # noqa: E501
+ assert result == expected
+
+
+def test_void_elements():
+ test_cases = [
+ ('', ''),
+ ('', ''),
+ ("Hello & Welcome
+ +Article 1
+Some text with bold and emphasis
+
+ ', ''),
+ ]
+
+ for input_html, expected in test_cases:
+ result, _ = set_html_attributes(input_html, ["data-root"], ["data-v-123"])
+ assert result == expected
+
+
+def test_html_head_with_meta():
+ html = """
+
+
+
+
+
"""
+
+ result: str
+ captured: Dict[str, List[str]]
+ result, captured = set_html_attributes(html, ["data-root"], ["data-v-123"], watch_on_attribute="data-id")
+ expected = """
+ Regular element
+ Nested element +
+
+
+
"""
+ assert result == expected
+
+ # Verify attribute capturing
+ assert len(captured) == 3
+
+ # Root element should have both root and all attributes
+ assert "123" in captured
+ assert "data-root" in captured["123"]
+ assert "data-v-123" in captured["123"]
+
+ # Non-root elements should only have all attributes
+ assert "456" in captured
+ assert captured["456"] == ["data-v-123"]
+ assert "789" in captured
+ assert captured["789"] == ["data-v-123"]
+
+
+def test_whitespace_preservation():
+ html = """Regular element
+ Nested element +
+
+
"""
+
+ result, _ = set_html_attributes(html, ["data-root"], ["data-all"])
+ expected = """Hello World
+ Text with spaces +
+
"""
+ assert result == expected
Hello World
+ Text with spaces +