♻️ refactor: rustify ass writer write_comment (#325)

Author: SigureMo

packages/biliass/rust/src/lib.rs

@@ -36,5 +36,6 @@ fn biliass_pyo3(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(python::py_find_alternative_row, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_mark_comment_row, m)?)?;
     m.add_function(wrap_pyfunction!(python::py_write_head, m)?)?;
+    m.add_function(wrap_pyfunction!(python::py_write_comment, m)?)?;
     Ok(())
 }

packages/biliass/rust/src/python/mod.rs

@@ -9,5 +9,5 @@ pub use reader::{py_read_comments_from_protobuf, py_read_comments_from_xml};
 pub use writer::{
     py_ass_escape, py_convert_color, py_convert_flash_rotation, py_convert_timestamp,
     py_find_alternative_row, py_get_zoom_factor, py_mark_comment_row, py_test_free_rows,
-    py_write_head, PyRows,
+    py_write_comment, py_write_head, PyRows,
 };

packages/biliass/rust/src/python/writer.rs

@@ -1,4 +1,5 @@
 use crate::python;
+use crate::writer::ass::write_comment;
 use crate::writer::{self, rows};
 
 use pyo3::prelude::*;
@@ -123,8 +124,8 @@ pub fn py_mark_comment_row(
 
 #[pyfunction(name = "write_head")]
 pub fn py_write_head(
-    width: i32,
-    height: i32,
+    width: u32,
+    height: u32,
     fontface: &str,
     fontsize: f32,
     alpha: f32,
@@ -134,3 +135,29 @@ pub fn py_write_head(
         width, height, fontface, fontsize, alpha, styleid,
     ))
 }
+
+#[allow(clippy::too_many_arguments)]
+#[pyfunction(name = "write_comment")]
+pub fn py_write_comment(
+    comment: &crate::python::PyComment,
+    row: usize,
+    width: u32,
+    height: u32,
+    bottom_reserved: u32,
+    fontsize: f32,
+    duration_marquee: f32,
+    duration_still: f32,
+    styleid: &str,
+) -> PyResult<String> {
+    Ok(write_comment(
+        &comment.inner,
+        row,
+        width,
+        height,
+        bottom_reserved,
+        fontsize,
+        duration_marquee,
+        duration_still,
+        styleid,
+    ))
+}

packages/biliass/rust/src/writer/ass.rs

@@ -1,6 +1,9 @@
+use crate::comment::{Comment, CommentPosition};
+use crate::writer::utils;
+
 pub fn write_head(
-    width: i32,
-    height: i32,
+    width: u32,
+    height: u32,
     fontface: &str,
     fontsize: f32,
     alpha: f32,
@@ -31,3 +34,64 @@ Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
 "
     )
 }
+
+fn convert_type2(row: usize, height: u32, bottom_reserved: u32) -> usize {
+    height as usize - bottom_reserved as usize - row
+}
+
+#[allow(clippy::too_many_arguments)]
+pub fn write_comment(
+    comment: &Comment,
+    row: usize,
+    width: u32,
+    height: u32,
+    bottom_reserved: u32,
+    fontsize: f32,
+    duration_marquee: f32,
+    duration_still: f32,
+    styleid: &str,
+) -> String {
+    let text = utils::ass_escape(&comment.comment);
+    let (style, duration) = match comment.pos {
+        CommentPosition::Bottom => {
+            let halfwidth = width / 2;
+            (format!("\\an8\\pos({halfwidth}, {row})"), duration_still)
+        }
+        CommentPosition::Top => {
+            let halfwidth = width / 2;
+            let row = convert_type2(row, height, bottom_reserved);
+            (format!("\\an2\\pos({halfwidth}, {row})"), duration_still)
+        }
+        CommentPosition::Reversed => {
+            let neglen = -(comment.width.ceil()) as i32;
+            (
+                format!("\\move({neglen}, {row}, {width}, {row})"),
+                duration_marquee,
+            )
+        }
+        _ => {
+            let neglen = -(comment.width.ceil()) as i32;
+            (
+                format!("\\move({width}, {row}, {neglen}, {row})"),
+                duration_marquee,
+            )
+        }
+    };
+    let mut styles = vec![style];
+    if comment.size - fontsize <= -1. || comment.size - fontsize >= 1. {
+        styles.push(format!("\\fs{:.0}", comment.size));
+    }
+    if comment.color != 0xFFFFFF {
+        styles.push(format!(
+            "\\c&H{}&",
+            utils::convert_color(comment.color, None, None)
+        ));
+        if comment.color == 0x000000 {
+            styles.push("\\3c&HFFFFFF&".to_owned());
+        }
+    }
+    let start = utils::convert_timestamp(comment.timeline);
+    let end = utils::convert_timestamp(comment.timeline + duration as f64);
+    let styles = styles.join("");
+    format!("Dialogue: 2,{start},{end},{styleid},,0000,0000,0000,,{{{styles}}}{text}\n")
+}

packages/biliass/src/biliass/_core.pyi

@@ -92,3 +92,14 @@ class Rows:
 def find_alternative_row(rows: Rows, comment: Comment, height: int, bottom_reserved: int) -> int: ...
 def mark_comment_row(rows: Rows, comment: Comment, row: int) -> None: ...
 def write_head(width: int, height: int, fontface: str, fontsize: float, alpha: float, styleid: str) -> str: ...
+def write_comment(
+    comment: Comment,
+    row: int,
+    width: int,
+    height: int,
+    bottom_reserved: int,
+    fontsize: float,
+    duration_marquee: float,
+    duration_still: float,
+    styleid: str,
+) -> str: ...

packages/biliass/src/biliass/biliass.py

@@ -23,6 +23,7 @@
     read_comments_from_protobuf,
     read_comments_from_xml,
     test_free_rows,
+    write_comment,
     write_head,
 )
 
@@ -160,44 +161,16 @@ def write_comment(
         duration_still: float,
         styleid: str,
     ):
-        text = ass_escape(comment.comment)
-        styles = []
-        if comment.pos == CommentPosition.Bottom:
-            styles.append("\\an8\\pos(%(halfwidth)d, %(row)d)" % {"halfwidth": width / 2, "row": row})
-            duration = duration_still
-        elif comment.pos == CommentPosition.Top:
-            styles.append(
-                "\\an2\\pos(%(halfwidth)d, %(row)d)"
-                % {
-                    "halfwidth": width / 2,
-                    "row": convert_type2(row, height, bottom_reserved),
-                }
-            )
-            duration = duration_still
-        elif comment.pos == CommentPosition.Reversed:
-            styles.append(
-                "\\move(%(neglen)d, %(row)d, %(width)d, %(row)d)"
-                % {"width": width, "row": row, "neglen": -math.ceil(comment.width)}
-            )
-            duration = duration_marquee
-        else:
-            styles.append(
-                "\\move(%(width)d, %(row)d, %(neglen)d, %(row)d)"
-                % {"width": width, "row": row, "neglen": -math.ceil(comment.width)}
-            )
-            duration = duration_marquee
-        if not (-1 < comment.size - fontsize < 1):
-            styles.append(f"\\fs{comment.size:.0f}")
-        if comment.color != 0xFFFFFF:
-            styles.append(f"\\c&H{convert_color(comment.color)}&")
-            if comment.color == 0x000000:
-                styles.append("\\3c&HFFFFFF&")
-        self._text += "Dialogue: 2,{start},{end},{styleid},,0000,0000,0000,,{{{styles}}}{text}\n".format(
-            start=convert_timestamp(comment.timeline),
-            end=convert_timestamp(comment.timeline + duration),
-            styles="".join(styles),
-            text=text,
-            styleid=styleid,
+        self._text += write_comment(
+            comment,
+            row,
+            width,
+            height,
+            bottom_reserved,
+            fontsize,
+            duration_marquee,
+            duration_still,
+            styleid,
         )
 
     def to_string(self):