refactor(webvtt): make WebVTTTimestamp public

Since WebVTTTimestamp is used in DoclingDocument, the class should be public.
Strengthen validation of cue language start tag annotation.

Signed-off-by: Cesar Berrospi Ramis <[email protected]>

Author: ceberam

docling_core/types/doc/document.py

@@ -68,7 +68,7 @@
 )
 from docling_core.types.doc.tokens import DocumentToken, TableToken
 from docling_core.types.doc.utils import parse_otsl_table_content, relative_path
-from docling_core.types.doc.webvtt import _WebVTTTimestamp
+from docling_core.types.doc.webvtt import WebVTTTimestamp
 
 _logger = logging.getLogger(__name__)
 
@@ -1230,14 +1230,14 @@ class ProvenanceTrack(BaseModel):
     """
 
     start_time: Annotated[
-        _WebVTTTimestamp,
+        WebVTTTimestamp,
         Field(
             examples=["00.11.000", "00:00:06.500", "01:28:34.300"],
             description="Start time offset of the track cue",
         ),
     ]
     end_time: Annotated[
-        _WebVTTTimestamp,
+        WebVTTTimestamp,
         Field(
             examples=["00.12.000", "00:00:08.200", "01:29:30.100"],
             description="End time offset of the track cue",

docling_core/types/doc/webvtt.py

@@ -28,9 +28,18 @@ class _WebVTTLineTerminator(str, Enum):
 ]
 
 
-class _WebVTTTimestamp(BaseModel):
+class WebVTTTimestamp(BaseModel):
     """WebVTT timestamp.
 
+    The timestamp is a string consisting of the following components in the given order:
+
+    - hours (optional, required if non-zero): two or more digits
+    - minutes: two digits between 0 and 59
+    - a colon character (:)
+    - seconds: two digits between 0 and 59
+    - a full stop character (.)
+    - thousandths of a second: three digits
+
     A WebVTT timestamp is always interpreted relative to the current playback position
     of the media data that the WebVTT file is to be synchronized with.
     """
@@ -54,6 +63,7 @@ class _WebVTTTimestamp(BaseModel):
 
     @model_validator(mode="after")
     def validate_raw(self) -> Self:
+        """Validate the WebVTT timestamp as a string."""
         m = self._pattern.match(self.raw)
         if not m:
             raise ValueError(f"Invalid WebVTT timestamp format: {self.raw}")
@@ -81,16 +91,15 @@ def seconds(self) -> float:
 
     @override
     def __str__(self) -> str:
+        """Return a string representation of a WebVTT timestamp."""
         return self.raw
 
 
 class _WebVTTCueTimings(BaseModel):
     """WebVTT cue timings."""
 
-    start: Annotated[
-        _WebVTTTimestamp, Field(description="Start time offset of the cue")
-    ]
-    end: Annotated[_WebVTTTimestamp, Field(description="End time offset of the cue")]
+    start: Annotated[WebVTTTimestamp, Field(description="Start time offset of the cue")]
+    end: Annotated[WebVTTTimestamp, Field(description="End time offset of the cue")]
 
     @model_validator(mode="after")
     def check_order(self) -> Self:
@@ -224,6 +233,21 @@ def __str__(self):
         return f"<{self._get_name_with_classes()} {self.annotation}>"
 
 
+class _WebVTTCueLanguageSpanStartTag(_WebVTTCueSpanStartTagAnnotated):
+    _bcp47_regex = re.compile(r"^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8})*$", re.IGNORECASE)
+
+    name: Literal["lang"] = Field("lang", description="The tag name")
+    annotation: Annotated[
+        str,
+        Field(
+            pattern=_bcp47_regex.pattern,
+            min_length=2,
+            max_length=99,
+            description="Cue language span start tag annotation",
+        ),
+    ]
+
+
 class _WebVTTCueComponentBase(BaseModel):
     """WebVTT caption or subtitle cue component.
 
@@ -294,7 +318,7 @@ class _WebVTTCueLanguageSpan(_WebVTTCueComponentBase):
     """
 
     kind: Literal["lang"] = "lang"
-    start_tag: _WebVTTCueSpanStartTagAnnotated
+    start_tag: _WebVTTCueLanguageSpanStartTag
 
 
 _WebVTTCueComponent = Annotated[
@@ -369,7 +393,7 @@ def parse(cls, raw: str) -> "_WebVTTCueBlock":
         start, end = [t.strip() for t in timing_line.split("-->")]
         end = re.split(" |\t", end)[0]  # ignore the cue settings list
         timings: _WebVTTCueTimings = _WebVTTCueTimings(
-            start=_WebVTTTimestamp(raw=start), end=_WebVTTTimestamp(raw=end)
+            start=WebVTTTimestamp(raw=start), end=WebVTTTimestamp(raw=end)
         )
         cue_text = " ".join(cue_lines).strip()
         # adding close tag for cue spans without end tag
@@ -409,13 +433,17 @@ def parse(cls, raw: str) -> "_WebVTTCueBlock":
                         classes: list[str] = []
                         if class_string:
                             classes = [c for c in class_string.split(".") if c]
-                        st = (
-                            _WebVTTCueSpanStartTagAnnotated(
+                        st: _WebVTTCueSpanStartTag
+                        if annotation and ct == "lang":
+                            st = _WebVTTCueLanguageSpanStartTag(
                                 name=ct, classes=classes, annotation=annotation.strip()
                             )
-                            if annotation
-                            else _WebVTTCueSpanStartTag(name=ct, classes=classes)
-                        )
+                        elif annotation:
+                            st = _WebVTTCueSpanStartTagAnnotated(
+                                name=ct, classes=classes, annotation=annotation.strip()
+                            )
+                        else:
+                            st = _WebVTTCueSpanStartTag(name=ct, classes=classes)
                         it = _WebVTTCueInternalText(components=children)
                         cp: _WebVTTCueComponent
                         if ct == "c":

docs/DoclingDocument.json

@@ -2143,7 +2143,7 @@
       "description": "Provenance information for elements extracted from media assets.\n\nA `ProvenanceTrack` instance describes a cue in a text track associated with a\nmedia element (audio, video, subtitles, screen recordings, ...).",
       "properties": {
         "start_time": {
-          "$ref": "#/$defs/_WebVTTTimestamp",
+          "$ref": "#/$defs/WebVTTTimestamp",
           "description": "Start time offset of the track cue",
           "examples": [
             "00.11.000",
@@ -2152,7 +2152,7 @@
           ]
         },
         "end_time": {
-          "$ref": "#/$defs/_WebVTTTimestamp",
+          "$ref": "#/$defs/WebVTTTimestamp",
           "description": "End time offset of the track cue",
           "examples": [
             "00.12.000",
@@ -3067,8 +3067,8 @@
       "title": "TitleItem",
       "type": "object"
     },
-    "_WebVTTTimestamp": {
-      "description": "WebVTT timestamp.\n\nA WebVTT timestamp is always interpreted relative to the current playback position\nof the media data that the WebVTT file is to be synchronized with.",
+    "WebVTTTimestamp": {
+      "description": "WebVTT timestamp.\n\nThe timestamp is a string consisting of the following components in the given order:\n\n- hours (optional, required if non-zero): two or more digits\n- minutes: two digits between 0 and 59\n- a colon character (:)\n- seconds: two digits between 0 and 59\n- a full stop character (.)\n- thousandths of a second: three digits\n\nA WebVTT timestamp is always interpreted relative to the current playback position\nof the media data that the WebVTT file is to be synchronized with.",
       "properties": {
         "raw": {
           "description": "A representation of the WebVTT Timestamp as a single string",
@@ -3079,7 +3079,7 @@
       "required": [
         "raw"
       ],
-      "title": "_WebVTTTimestamp",
+      "title": "WebVTTTimestamp",
       "type": "object"
     }
   },

test/test_webvtt.py

@@ -9,17 +9,18 @@
 from pydantic import ValidationError
 
 from docling_core.types.doc.webvtt import (
+    WebVTTTimestamp,
     _WebVTTCueBlock,
     _WebVTTCueComponentWithTerminator,
     _WebVTTCueInternalText,
     _WebVTTCueItalicSpan,
     _WebVTTCueLanguageSpan,
+    _WebVTTCueLanguageSpanStartTag,
     _WebVTTCueSpanStartTagAnnotated,
     _WebVTTCueTextSpan,
     _WebVTTCueTimings,
     _WebVTTCueVoiceSpan,
     _WebVTTFile,
-    _WebVTTTimestamp,
 )
 
 from .test_data_gen_flag import GEN_TEST_DATA
@@ -42,7 +43,7 @@ def test_vtt_cue_commponents() -> None:
         0.0,
     ]
     for idx, ts in enumerate(valid_timestamps):
-        model = _WebVTTTimestamp(raw=ts)
+        model = WebVTTTimestamp(raw=ts)
         assert model.seconds == valid_total_seconds[idx]
 
     """Test invalid WebVTT timestamps."""
@@ -57,35 +58,35 @@ def test_vtt_cue_commponents() -> None:
     ]
     for ts in invalid_timestamps:
         with pytest.raises(ValidationError):
-            _WebVTTTimestamp(raw=ts)
+            WebVTTTimestamp(raw=ts)
 
     """Test the timestamp __str__ method."""
-    model = _WebVTTTimestamp(raw="00:01:02.345")
+    model = WebVTTTimestamp(raw="00:01:02.345")
     assert str(model) == "00:01:02.345"
 
     """Test valid cue timings."""
-    start = _WebVTTTimestamp(raw="00:10.005")
-    end = _WebVTTTimestamp(raw="00:14.007")
+    start = WebVTTTimestamp(raw="00:10.005")
+    end = WebVTTTimestamp(raw="00:14.007")
     cue_timings = _WebVTTCueTimings(start=start, end=end)
     assert cue_timings.start == start
     assert cue_timings.end == end
     assert str(cue_timings) == "00:10.005 --> 00:14.007"
 
     """Test invalid cue timings with end timestamp before start."""
-    start = _WebVTTTimestamp(raw="00:10.700")
-    end = _WebVTTTimestamp(raw="00:10.500")
+    start = WebVTTTimestamp(raw="00:10.700")
+    end = WebVTTTimestamp(raw="00:10.500")
     with pytest.raises(ValidationError) as excinfo:
         _WebVTTCueTimings(start=start, end=end)
     assert "End timestamp must be greater than start timestamp" in str(excinfo.value)
 
     """Test invalid cue timings with missing end."""
-    start = _WebVTTTimestamp(raw="00:10.500")
+    start = WebVTTTimestamp(raw="00:10.500")
     with pytest.raises(ValidationError) as excinfo:
         _WebVTTCueTimings(start=start)  # type: ignore[call-arg]
     assert "Field required" in str(excinfo.value)
 
     """Test invalid cue timings with missing start."""
-    end = _WebVTTTimestamp(raw="00:10.500")
+    end = WebVTTTimestamp(raw="00:10.500")
     with pytest.raises(ValidationError) as excinfo:
         _WebVTTCueTimings(end=end)  # type: ignore[call-arg]
     assert "Field required" in str(excinfo.value)
@@ -272,3 +273,13 @@ def test_webvtt_file() -> None:
     assert len(block.payload) == 1
     assert isinstance(block.payload[0].component, _WebVTTCueTextSpan)
     assert block.payload[0].component.text == "Good."
+
+
+def test_webvtt_cue_language_span_start_tag():
+    _WebVTTCueLanguageSpanStartTag.model_validate_json('{"annotation": "en"}')
+    _WebVTTCueLanguageSpanStartTag.model_validate_json('{"annotation": "en-US"}')
+    _WebVTTCueLanguageSpanStartTag.model_validate_json('{"annotation": "zh-Hant"}')
+    with pytest.raises(ValidationError, match="should match pattern"):
+        _WebVTTCueLanguageSpanStartTag.model_validate_json('{"annotation": "en_US"}')
+    with pytest.raises(ValidationError, match="should match pattern"):
+        _WebVTTCueLanguageSpanStartTag.model_validate_json('{"annotation": "123-de"}')