feat: support reading existing chapters from MKV files

pymkv/MKVFile.py

@@ -62,6 +62,7 @@
 import subprocess as sp
 import sys
 import tempfile
+import xml.etree.ElementTree as ET
 from collections.abc import Callable, Iterable, Sequence
 from pathlib import Path
 from typing import Any, ClassVar, TypeVar, cast
@@ -79,6 +80,7 @@
     Chapters,
     EditionEntry,
     export_to_xml,
+    parse_chapters_xml,
 )
 from pymkv.command_generators import (
     AttachmentOptions,
@@ -135,6 +137,10 @@ class MKVFile:
         The path where pymkv looks for the mkvmerge executable. pymkv relies on the mkvmerge executable to parse
         files. By default, it is assumed mkvmerge is in your shell's $PATH variable. If it is not, you need to set
         *mkvmerge_path* to the executable location.
+    mkvextract_path : str, optional
+        The path where pymkv looks for the mkvextract executable. By default, it is
+        assumed mkvextract is in your shell's $PATH variable. If it is not, you need to set
+        *mkvextract_path* to the executable location.
 
     Raises
     ------
@@ -149,11 +155,14 @@ def __init__(
         file_path: str | os.PathLike | None = None,
         title: str | None = None,
         mkvmerge_path: str | os.PathLike | Iterable[str] = "mkvmerge",
+        mkvextract_path: str | os.PathLike | Iterable[str] = "mkvextract",
     ) -> None:
         self.mkvmerge_path: tuple[str, ...] = prepare_mkvtoolnix_path(mkvmerge_path)
+        self.mkvextract_path: tuple[str, ...] = prepare_mkvtoolnix_path(mkvextract_path)
         self.title = title
         self._chapters_file: str | None = None
-        self.chapters_obj: Chapters | None = None
+        self._chapters_obj: Chapters | None = None
+        self._chapters_from_source: bool = False
         self._temp_chapters_file: str | None = None
         self._chapter_language: str | None = None
         self._global_tags_file: str | None = None
@@ -248,6 +257,13 @@ def __init__(
                     new_attachment.source_file = file_path
                     self.attachments.append(new_attachment)
 
+            chapter_entries = sum(c.num_entries for c in info_struct.chapters)
+            if chapter_entries > 0:
+                chapters_from_source = self._read_chapters(file_path)
+                if chapters_from_source is not None:
+                    self._chapters_obj = chapters_from_source
+                    self._chapters_from_source = True
+
         # split options
         self._split_options: list[str] = []
         self._progress_handler = None
@@ -280,6 +296,28 @@ def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
         """Clean up temporary files on context manager exit."""
         self.cleanup()
 
+    @property
+    def chapters_obj(self) -> Chapters | None:
+        """
+        Get or set the :class:`~pymkv.chapters.Chapters` object attached to this file.
+
+        Assigning to this property (``mkv.chapters_obj = Chapters()``) always marks the
+        chapters as caller-supplied: :meth:`command` will serialize them to a temporary XML
+        file and pass it to mkvmerge via ``--chapters``, overriding whatever chapters (if any)
+        the source file already had.
+
+        Returns
+        -------
+        Chapters | None
+            The chapters object, or ``None`` if no chapters are set.
+        """
+        return self._chapters_obj
+
+    @chapters_obj.setter
+    def chapters_obj(self, chapters: Chapters | None) -> None:
+        self._chapters_obj = chapters
+        self._chapters_from_source = False
+
     @property
     def chapter_language(self) -> str | None:
         """
@@ -390,7 +428,7 @@ def command(
         self.output_path = str(Path(output_path).expanduser())
 
         # Handle object-based chapters
-        if self.chapters_obj and not self._chapters_file:
+        if self.chapters_obj and not self._chapters_file and not self._chapters_from_source:
             self._write_chapters_xml()
 
         # Pre-assign file IDs
@@ -528,6 +566,42 @@ def mux(
 
         return proc.returncode
 
+    def _read_chapters(self, file_path: str) -> Chapters | None:
+        """
+        Extract and parse the chapters already present in `file_path`.
+
+        Uses ``mkvextract chapters <file_path>`` to obtain the chapter XML and parses it into a
+        :class:`~pymkv.chapters.Chapters` object. This is used internally when importing a
+        pre-existing MKV file that has chapters. Failures are logged and treated as "no chapters"
+        rather than raised, so a file with unreadable chapters can still be opened.
+
+        Parameters
+        ----------
+        file_path : str
+            The path of the MKV file to read chapters from.
+
+        Returns
+        -------
+        Chapters | None
+            The parsed chapters, or ``None`` if they could not be extracted or parsed.
+        """
+        command = [*self.mkvextract_path, "chapters", file_path]
+        try:
+            result = sp.run(command, check=True, capture_output=True)  # noqa: S603
+        except (sp.CalledProcessError, FileNotFoundError) as e:
+            logging.warning("Could not extract chapters from '%s': %s", file_path, e)
+            return None
+
+        xml_content = result.stdout.decode("utf-8")
+        if not xml_content.strip():
+            return None
+
+        try:
+            return parse_chapters_xml(xml_content)
+        except ET.ParseError as e:
+            logging.warning("Could not parse chapters XML from '%s': %s", file_path, e)
+            return None
+
     def _write_chapters_xml(self) -> None:
         """
         Write chapter objects to a temporary XML file.
@@ -580,6 +654,8 @@ def add_chapter(self, chapter: ChapterAtom | EditionEntry) -> None:
         """
         if self.chapters_obj is None:
             self.chapters_obj = Chapters()
+        else:
+            self._chapters_from_source = False
 
         if isinstance(chapter, ChapterAtom):
             if not self.chapters_obj.editions:

pymkv/chapters.py

@@ -161,3 +161,133 @@ def export_to_xml(chapters: Chapters) -> str:
     xml_bytes = ET.tostring(root, encoding="utf-8", xml_declaration=True)
 
     return xml_bytes.decode("utf-8")
+
+
+def _text_to_bool(value: str | None) -> bool | None:
+    """
+    Convert a Matroska chapter XML flag ("0"/"1") to a bool, preserving ``None``.
+    """
+    if value is None:
+        return None
+    return value.strip() == "1"
+
+
+def _text_to_int(value: str | None) -> int | None:
+    """
+    Convert element text to an int, preserving ``None`` and tolerating blank/invalid values.
+    """
+    if value is None:
+        return None
+    value = value.strip()
+    if not value:
+        return None
+    try:
+        return int(value)
+    except ValueError:
+        return None
+
+
+def _findtext(element: ET.Element, tag: str) -> str | None:
+    """
+    Return the stripped text of a child element, or ``None`` if it is absent or blank.
+    """
+    text = element.findtext(tag)
+    if text is None:
+        return None
+    text = text.strip()
+    return text or None
+
+
+def _strip_none(data: dict[str, Any]) -> dict[str, Any]:
+    """
+    Drop keys whose value is ``None`` before handing the dict to :func:`msgspec.convert`.
+    """
+    return {key: value for key, value in data.items() if value is not None}
+
+
+def _parse_chapter_display(element: ET.Element) -> dict[str, Any]:
+    """
+    Parse a single ``<ChapterDisplay>`` element into a dict matching :class:`ChapterDisplay`.
+    """
+    return _strip_none(
+        {
+            "ChapterString": _findtext(element, "ChapterString") or "",
+            "ChapterLanguage": _findtext(element, "ChapterLanguage"),
+            "ChapterCountry": _findtext(element, "ChapterCountry"),
+        }
+    )
+
+
+def _parse_chapter_atom(element: ET.Element) -> dict[str, Any]:
+    """
+    Recursively parse a ``<ChapterAtom>`` element into a dict matching :class:`ChapterAtom`.
+    """
+    return _strip_none(
+        {
+            "ChapterTimeStart": _findtext(element, "ChapterTimeStart") or "",
+            "ChapterTimeEnd": _findtext(element, "ChapterTimeEnd"),
+            "ChapterUID": _text_to_int(element.findtext("ChapterUID")),
+            "ChapterFlagHidden": _text_to_bool(element.findtext("ChapterFlagHidden")),
+            "ChapterFlagEnabled": _text_to_bool(element.findtext("ChapterFlagEnabled")),
+            "ChapterDisplay": [_parse_chapter_display(d) for d in element.findall("ChapterDisplay")],
+            "ChapterAtom": [_parse_chapter_atom(a) for a in element.findall("ChapterAtom")],
+        }
+    )
+
+
+def _parse_edition_entry(element: ET.Element) -> dict[str, Any]:
+    """
+    Parse a single ``<EditionEntry>`` element into a dict matching :class:`EditionEntry`.
+    """
+    return _strip_none(
+        {
+            "EditionUID": _text_to_int(element.findtext("EditionUID")),
+            "EditionFlagHidden": _text_to_bool(element.findtext("EditionFlagHidden")),
+            "EditionFlagDefault": _text_to_bool(element.findtext("EditionFlagDefault")),
+            "EditionFlagOrdered": _text_to_bool(element.findtext("EditionFlagOrdered")),
+            "ChapterAtom": [_parse_chapter_atom(a) for a in element.findall("ChapterAtom")],
+        }
+    )
+
+
+def parse_chapters_xml(xml_content: str | bytes) -> Chapters:
+    """
+    Parse a Matroska chapters XML document (such as the output of ``mkvextract chapters``)
+    into a :class:`Chapters` object.
+
+    Parameters
+    ----------
+    xml_content : str | bytes
+        The XML content to parse, as returned by ``mkvextract chapters <file>``.
+
+    Returns
+    -------
+    Chapters
+        The parsed chapters, with one :class:`EditionEntry` per ``<EditionEntry>`` element.
+
+    Raises
+    ------
+    xml.etree.ElementTree.ParseError
+        If `xml_content` is not well-formed XML.
+
+    Examples
+    --------
+    >>> xml = '''<?xml version="1.0"?>
+    ... <Chapters>
+    ...   <EditionEntry>
+    ...     <ChapterAtom>
+    ...       <ChapterTimeStart>00:00:00.000</ChapterTimeStart>
+    ...       <ChapterDisplay>
+    ...         <ChapterString>Intro</ChapterString>
+    ...         <ChapterLanguage>eng</ChapterLanguage>
+    ...       </ChapterDisplay>
+    ...     </ChapterAtom>
+    ...   </EditionEntry>
+    ... </Chapters>'''
+    >>> chapters = parse_chapters_xml(xml)
+    >>> chapters.editions[0].atoms[0].displays[0].string
+    'Intro'
+    """
+    root = ET.fromstring(xml_content)  # noqa: S314
+    data = {"EditionEntry": [_parse_edition_entry(e) for e in root.findall("EditionEntry")]}
+    return msgspec.convert(data, type=Chapters, strict=False)

pymkv/models.py

@@ -168,6 +168,19 @@ class AttachmentInfo(msgspec.Struct):
     properties: AttachmentProperties = msgspec.field(default_factory=AttachmentProperties)
 
 
+class ChapterInfo(msgspec.Struct):
+    """
+    Summary information about the chapters present in a file.
+
+    Attributes
+    ----------
+    num_entries : int
+        The number of chapter entries.
+    """
+
+    num_entries: int = 0
+
+
 class MkvMergeOutput(msgspec.Struct):
     """
     Root structure of `mkvmerge -J` output.
@@ -184,6 +197,8 @@ class MkvMergeOutput(msgspec.Struct):
         List of track tags.
     attachments : list[AttachmentInfo]
         List of attachments.
+    chapters : list[ChapterInfo]
+        Summary of chapters present in the file.
     file_name : str | None
         The file name.
     """
@@ -194,3 +209,4 @@ class MkvMergeOutput(msgspec.Struct):
     global_tags: list[TagEntry] = []
     track_tags: list[TagEntry] = []
     attachments: list[AttachmentInfo] = []
+    chapters: list[ChapterInfo] = []