Overhaul frame props plugin

vspreview/plugins/builtins/frame_props.ppy

@@ -1,128 +1,239 @@
 from __future__ import annotations
 
-from PyQt6.QtWidgets import QLabel, QSizePolicy, QTableView, QWidget
-from vapoursynth import FrameProps
-from vstools import ChromaLocation, ColorRange, FieldBased, Matrix, Primaries, PropEnum, Transfer
+from typing import Any
 
+from PyQt6.QtCore import QModelIndex
+from PyQt6.QtWidgets import QApplication, QHeaderView, QLabel, QTableView, QWidget
+from vapoursynth import FrameProps
 from vspreview.core import Frame, HBoxLayout, Stretch, Switch, TableModel, VBoxLayout
 from vspreview.plugins import AbstractPlugin, PluginConfig
 
+from .frame_props.category import (
+    frame_props_categories, frame_props_category_prefix_lut, frame_props_category_suffix_lut
+)
+from .frame_props.exclude import frame_props_excluded_keys
+from .frame_props.lut import frame_props_lut
+
+
 __all__ = [
     'FramePropsPlugin'
 ]
 
 
-def _create_enum_props_lut(enum: type[PropEnum], pretty_name: str) -> tuple[str, dict[str, dict[int, str]]]:
-    return enum.prop_key, {
-        pretty_name: {
-            idx: enum.from_param(idx).pretty_string if enum.is_valid(idx) else 'Invalid'
-            for idx in range(min(enum.__members__.values()) - 1, max(enum.__members__.values()) + 1)
-        }
-    }
-
-
-_frame_props_excluded_keys = {
-    # vs internals
-    '_AbsoluteTime', '_DurationNum', '_DurationDen', '_PictType', '_Alpha',
-    # Handled separately
-    '_SARNum', '_SARDen',
-    # source filters
-    '_FrameNumber',
-    # vstools set_output
-    'Name'
-}
-
-
-_frame_props_lut = {
-    '_Combed': {
-        'Is Combed': [
-            'No',
-            'Yes'
-        ]
-    },
-    '_Field': {
-        'Frame Field Type': [
-            'Bottom Field',
-            'Top Field'
-        ]
-    },
-    '_SceneChangeNext': {
-        'Scene Cut': [
-            'Current Scene',
-            'End of Scene'
-        ]
-    },
-    '_SceneChangePrev': {
-        'Scene Change': [
-            'Current Scene',
-            'Start of Scene'
-        ]
-    }
-} | dict([
-    _create_enum_props_lut(enum, name)
-    for enum, name in list[tuple[type[PropEnum], str]]([
-        (FieldBased, 'Field Type'),
-        (Matrix, 'Matrix'),
-        (Transfer, 'Transfer'),
-        (Primaries, 'Primaries'),
-        (ChromaLocation, 'Chroma Location'),
-        (ColorRange, 'Color Range')
-    ])
-])
+class FramePropsPlugin(AbstractPlugin, QWidget):
+    """
+    A plugin for displaying and managing frame properties in VSPreview.
 
+    This plugin provides a user interface for viewing frame properties of video clips.
+    It supports both prettified and raw data views of frame properties,
+    organized into categories for easy navigation.
+    """
 
-class FramePropsPlugin(AbstractPlugin, QWidget):
     _config = PluginConfig('dev.setsugen.frame_props', 'Frame Props')
 
     def setup_ui(self) -> None:
-        self.table = QTableView()
-        self.table._model = TableModel([], ['Name', 'Data'], False)
-        self.table.verticalHeader().hide()
-        self.table.horizontalHeader().setStretchLastSection(True)
-        self.table.horizontalHeader().setFirstSectionMovable(False)
-
-        def _update_prettify(clicked: bool) -> None:
-            self.settings.local.prettify = clicked
+        """Set up the user interface components."""
+
+        self.setup_tables()
+        self.setup_raw_data_switch()
+        self.setup_labels()
+        self.setup_layout()
+
+    def setup_tables(self) -> None:
+        """Initialize and set up the table views."""
+
+        self.category_tables = {category: self.create_table_view(category) for category in frame_props_categories}
+        self.raw_data_table = self.create_table_view('raw_data')
+
+        self.raw_data_table.hide()
+
+    def create_table_view(self, category: str) -> QTableView:
+        """Create and configure a QTableView."""
+
+        table = QTableView()
+        table._model = TableModel([], ['Name', 'Data'], False)
+
+        table.verticalHeader().hide()
+        table.horizontalHeader().setStretchLastSection(True)
+        table.horizontalHeader().setFirstSectionMovable(False)
+
+        if category != 'Other':
+            table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.ResizeToContents)
+            table.verticalHeader().setSectionResizeMode(QHeaderView.ResizeMode.ResizeToContents)
+        else:
+            table.horizontalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Interactive)
+            table.verticalHeader().setSectionResizeMode(QHeaderView.ResizeMode.Interactive)
+
+        table.doubleClicked.connect(self.copy_value_to_clipboard)
+
+        return table
+
+    def setup_raw_data_switch(self) -> None:
+        """Set up the raw data switch."""
+
+        def _update_raw_data(clicked: bool) -> None:
+            self.settings.local.raw_data = clicked
             self.on_current_frame_changed(None)
 
-        self.prettify = Switch(10, 24, clicked=_update_prettify, tooltip='Pretty')
-        if 'prettify' not in self.settings.local or self.settings.local.prettify:
-            self.prettify.click()
+        self.raw_data = Switch(10, 24, clicked=_update_raw_data, tooltip='Display the raw frame properties')
+
+        if 'raw_data' not in self.settings.local:
+            self.settings.local.raw_data = False
+
+    def setup_labels(self) -> None:
+        """Set up the labels for different sections."""
+
+        self.category_labels = {category: QLabel(f'{category} Properties:') for category in frame_props_categories}
+
+    def setup_layout(self) -> None:
+        """Set up the layout of the widget."""
+
+        layout_items = [HBoxLayout([QLabel('Raw data:'), self.raw_data, Stretch()])]
 
-        VBoxLayout(self, [
-            HBoxLayout([QLabel('Prettify output:'), self.prettify, Stretch()]),
-            self.table
+        for category in frame_props_categories:
+            layout_items.extend([
+                self.category_labels[category],
+                self.category_tables[category]
+            ])
+
+        layout_items.extend([
+            self.raw_data_table
         ])
 
+        VBoxLayout(self, layout_items)
+
     def on_current_frame_changed(self, frame: Frame) -> None:
+        """Handle the event when the current frame changes."""
+
         if (props := self.main.current_output.props) is None:
             return
 
         self.update_frame_props(props)
 
     def update_frame_props(self, props: FrameProps) -> None:
-        self.table._model._data = []
+        """Update the frame properties display."""
+
+        self.clear_table_data()
+
+        if self.raw_data.isChecked():
+            self.populate_raw_table(props)
+        else:
+            self.populate_prettified_tables(props)
+
+        self.update_table_models()
+        self.toggle_table_visibility()
+
+    def clear_table_data(self) -> None:
+        """Clear the data in all tables."""
 
-        if self.prettify.isChecked():
-            for key in sorted(props.keys()):
-                if key in _frame_props_excluded_keys:
-                    continue
+        for table in list(self.category_tables.values()) + [self.raw_data_table]:
+            table._model._data = []
 
-                if key in _frame_props_lut:
-                    title = next(iter(_frame_props_lut[key].keys()))
-                    value_str = _frame_props_lut[key][title][props[key]]  # type: ignore
-                else:
-                    title = key[1:] if key.startswith('_') else key
-                    value_str = str(props[key])
+    def populate_raw_table(self, props: FrameProps) -> None:
+        """Populate the raw data table with all frame properties."""
+
+        for key in self.sort_props(props):
+            value = props[key]
+            self.raw_data_table._model._data.append([key, str(value)])
+
+    def populate_prettified_tables(self, props: FrameProps) -> None:
+        """Populate tables with prettified frame properties."""
+
+        stripped_props = {
+            k.strip(): v.strip() if isinstance(v, str) else v
+            for k, v in props.items()
+        }
+
+        prettified_props = {}
+
+        for key, value in stripped_props.items():
+            if key not in frame_props_excluded_keys:
+                title, value_str = self.get_prettified_prop(key, value)
 
                 if value_str is not None:
-                    self.table._model._data.append([title, value_str])
+                    prettified_props[key] = (title, value_str)
+
+        for key in self.sort_props(prettified_props):
+            title, value_str = prettified_props[key]
+            category = self.get_property_category(key)
+            self.category_tables[category]._model._data.append([title, value_str])
+
+        self.add_pixel_aspect_ratio(stripped_props)
 
-            if '_SARNum' in props and '_SARDen' in props:
-                self.table._model._data.append(['Pixel aspect ratio', f"{props['_SARNum']}/{props['_SARDen']}"])
+    def get_property_category(self, key: str) -> str:
+        """Determine the category for a given property key."""
+
+        for prefix, category in frame_props_category_prefix_lut.items():
+            if key.startswith(prefix):
+                return str(category)
+
+        for suffix, category in frame_props_category_suffix_lut.items():
+            if key.endswith(suffix):
+                return str(category)
+
+        for category, props in frame_props_categories.items():
+            if key in props:
+                return str(category)
+
+        return 'Other'
+
+    def get_prettified_prop(self, key: str, value: Any) -> tuple[str, str | None]:
+        """Get a prettified version of a property."""
+
+        if key in frame_props_lut:
+            title = next(iter(frame_props_lut[key].keys()))
+
+            value_str = str(
+                frame_props_lut[key][title](value) if callable(frame_props_lut[key][title])
+                else frame_props_lut[key][title][value]
+            )
         else:
-            for key, value in props.items():
-                self.table._model._data.append([key, str(value)])
+            title = key[1:] if key.startswith('_') else key
+            value_str = str(value)
+
+        return title, value_str
+
+    def add_pixel_aspect_ratio(self, props: FrameProps) -> None:
+        """Add pixel aspect ratio to the video properties table if available."""
+
+        sar_num, sar_den = props.get('_SARNum'), props.get('_SARDen')
+
+        if sar_num is None or sar_den is None:
+            return
+
+        self.category_tables['Video']._model._data.append(
+            ['Sample aspect ratio', f"{sar_num}/{sar_den}"]
+        )
+
+    def sort_props(self, props: FrameProps) -> list[str]:
+        """Sort properties with underscore-prefixed keys first."""
+
+        return sorted(props.keys(), key=lambda x: (not x.startswith('_'), x))
+
+    def update_table_models(self) -> None:
+        """Update the models of all tables."""
+
+        for table in list(self.category_tables.values()) + [self.raw_data_table]:
+            table.setModel(table._model)
+            table._model.layoutChanged.emit()
+
+    def toggle_table_visibility(self) -> None:
+        """Toggle the visibility of tables based on raw data mode."""
+
+        is_raw = self.raw_data.isChecked()
+
+        for category, table in self.category_tables.items():
+            is_visible = not is_raw and len(table._model._data) > 0
+            table.setVisible(is_visible)
+
+            if hasattr(self, 'category_labels'):
+                self.category_labels[category].setVisible(is_visible)
+
+        self.raw_data_table.setVisible(is_raw)
+
+    def copy_value_to_clipboard(self, index: QModelIndex) -> None:
+        """Copy the value of the clicked cell to the clipboard."""
+
+        value = index.sibling(index.row(), 1).data()
 
-        self.table.setModel(self.table._model)
-        self.table._model.layoutChanged.emit()
+        QApplication.clipboard().setText(value)