ami/ami-columns/__init__.py

PLUGIN_NAME = "Ami (Columns)"
PLUGIN_AUTHOR = "SuperSaltyGamer"
PLUGIN_DESCRIPTION = "Custom computed columns in cluster and album panes."
PLUGIN_VERSION = "1.0.0"
PLUGIN_API_VERSIONS = ["2.10"]
PLUGIN_LICENSE = "AGPL-3.0"

from typing import Dict, Tuple
from picard import log
from picard.album import Album
from picard.config import ListOption, IntOption, config
from picard.file import File
from picard.metadata import Metadata
from picard.script import ScriptParser
from picard.ui.itemviews import MainPanel
from picard.ui.options import OptionsPage, register_options_page
from picard.ui.util import qlistwidget_items
from picard.ui.widgets.scriptlistwidget import ScriptListWidget, ScriptListWidgetItem
from picard.ui.widgets.scripttextedit import ScriptTextEdit
from PyQt5 import QtCore, QtWidgets

OPTION_COLUMNS = "ami_columns_columns"
OPTION_SELECTED_COLUMN = "ami_columns_selected_column"

computed_columns: Dict[str, Tuple[bool, str]] = {}
script_parser = ScriptParser()

original_file_column = File.column
original_album_column = Album.column


def override_file_column(self: File, column_name: str):
	try:
		column = computed_columns.get(column_name, None)
		if column is None or not column[0]:  # Handle built-in columns as usual.
			return original_file_column(self, column_name)
		new_metadata = Metadata()
		for name in self.metadata:
			new_metadata[name] = self.metadata.getall(name)
		value = script_parser.eval(column[1], new_metadata, self) or ""
		if not hasattr(self, "_computed_column_values"):
			self._computed_column_values = {}
		self._computed_column_values[column_name] = value
		return value
	except Exception:
		log.warning("Failed to compute column '%s' for %s", column_name, self, exc_info=True)


def override_album_column(self: Album, column_name: str):
	try:
		column = computed_columns.get(column_name, None)
		if column is None or not column[0]:  # Handle built-in columns as usual.
			return original_album_column(self, column_name)
		values = set()
		for track in self.tracks:
			for file in track.files:
				if hasattr(file, "_computed_column_values"):
					values.add(file._computed_column_values.get(column_name, ""))
		values_count = len(values)
		if values_count == 0:
			return ""
		elif values_count == 1:
			return values.pop()
		else:
			return f"(different across {len(values)} items)"
	except Exception:
		log.warning("Failed to compute column '%s' for %s", column_name, self, exc_info=True)


File.column = override_file_column
Album.column = override_album_column


def update_columns():
	global computed_columns
	computed_columns = {name: (enabled, text) for name, enabled, text in config.setting[OPTION_COLUMNS]}
	MainPanel.columns = [column for column in MainPanel.columns if column[0] != column[1]]
	for name in computed_columns.keys():
		MainPanel.columns.append((name, name))
	MainPanel._column_indexes = {column[1]: i for i, column in enumerate(MainPanel.columns)}


class PluginOptionsPage(OptionsPage):
	NAME = "ami_columns"
	TITLE = "Ami (Columns)"
	PARENT = "plugins"

	options = [
		ListOption("setting", OPTION_COLUMNS, []),
		IntOption("persist", OPTION_SELECTED_COLUMN, 0)
	]

	def __init__(self, parent=None):
		super(PluginOptionsPage, self).__init__(parent)
		self.ui = Ui_PluginOptionsPage(self)

	def load(self):
		self.ui.script_list.clear()
		for name, enabled, text in config.setting[OPTION_COLUMNS]:
			self.ui.script_list.addItem(ScriptListWidgetItem(name, enabled, text))
		item = self.ui.script_list.item(config.persist[OPTION_SELECTED_COLUMN])
		if item:
			self.ui.script_list.setCurrentItem(item)
			item.setSelected(True)

	def save(self):
		scripts: list[(str, bool, str)] = []
		for item in qlistwidget_items(self.ui.script_list):
			pos, name, enabled, text = item.get_all()
			scripts.append((name, enabled, text))
		config.setting[OPTION_COLUMNS] = scripts
		config.persist[OPTION_SELECTED_COLUMN] = self.ui.script_list.currentRow()
		update_columns()

	def add_script(self):
		item = ScriptListWidgetItem(f"Column")
		item.setCheckState(QtCore.Qt.CheckState.Checked)
		self.ui.script_list.addItem(item)
		self.ui.script_list.setCurrentItem(item, QtCore.QItemSelectionModel.SelectionFlag.Clear | QtCore.QItemSelectionModel.SelectionFlag.SelectCurrent)

	def remove_script(self):
		items = self.ui.script_list.selectedItems()
		if items:
			item = self.ui.script_list.takeItem(self.ui.script_list.row(items[0]))
			del item

	def script_selected(self):
		items = self.ui.script_list.selectedItems()
		if items:
			item = items[0]
			self.ui.script_edit.setEnabled(True)
			self.ui.script_edit.setText(item.script)
			self.ui.script_edit.setFocus(QtCore.Qt.FocusReason.OtherFocusReason)
		else:
			self.ui.script_edit.setEnabled(False)
			self.ui.script_edit.setText("")

	def live_update_and_check(self):
		items = self.ui.script_list.selectedItems()
		if not items:
			return
		item = items[0]
		item.script = self.ui.script_edit.toPlainText()
		self.ui.error_label.setStyleSheet("")
		self.ui.error_label.setText("")
		try:
			script_parser.eval(item.script)
			item.has_error = False
		except Exception as err:
			item.has_error = True
			self.ui.error_label.setStyleSheet(self.STYLESHEET_ERROR)
			self.ui.error_label.setText(str(err))
			return


class Ui_PluginOptionsPage(object):

	def __init__(self, page: PluginOptionsPage):
		page.setObjectName("plugin_options_page")
		self.vertical_layout = QtWidgets.QVBoxLayout(page)
		self.vertical_layout.setObjectName("vertical_layout")
		self.warn_label = QtWidgets.QLabel(page)
		self.warn_label.setObjectName("warn_label")
		self.warn_label.setText("<font color='#ff3333'><b>Restart Picard before modifying visible columns in panes.</b></font>")
		self.vertical_layout.addWidget(self.warn_label)
		self.info_label = QtWidgets.QLabel(page)
		self.info_label.setObjectName("info_label")
		self.info_label.setText("Enabled columns will be computed in cluster and album panes.")
		self.vertical_layout.addWidget(self.info_label)
		self.splitter = QtWidgets.QSplitter(page)
		self.splitter.setObjectName("splitter")
		self.splitter.setOrientation(QtCore.Qt.Orientation.Horizontal)
		size_policy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Policy.Expanding, QtWidgets.QSizePolicy.Policy.Expanding)
		size_policy.setHorizontalStretch(0)
		size_policy.setVerticalStretch(0)
		size_policy.setHeightForWidth(self.splitter.sizePolicy().hasHeightForWidth())
		self.splitter.setSizePolicy(size_policy)
		self.script_list = ScriptListWidget(self.splitter)
		self.script_list.setObjectName("script_list")
		self.script_list.setMinimumSize(QtCore.QSize(120, 0))
		self.script_list.itemSelectionChanged.connect(page.script_selected)
		self.form_widget = QtWidgets.QWidget(self.splitter)
		self.form_widget.setObjectName("form_widget")
		self.vertical_layout_2 = QtWidgets.QVBoxLayout(self.form_widget)
		self.vertical_layout_2.setContentsMargins(0, 0, 0, 0)
		self.vertical_layout_2.setObjectName("vertical_layout_2")
		self.script_edit = ScriptTextEdit(self.form_widget)
		self.script_edit.setObjectName("script_edit")
		self.script_edit.setAcceptRichText(False)
		self.script_edit.textChanged.connect(page.live_update_and_check)
		self.vertical_layout_2.addWidget(self.script_edit)
		self.vertical_layout.addWidget(self.splitter)
		self.horizontal_layout = QtWidgets.QHBoxLayout()
		self.horizontal_layout.setObjectName("horizontal_layout")
		self.add_button = QtWidgets.QToolButton(page)
		self.add_button.setObjectName("add_button")
		self.add_button.setText("Add new column")
		self.add_button.clicked.connect(page.add_script)
		self.horizontal_layout.addWidget(self.add_button)
		self.remove_button = QtWidgets.QToolButton(page)
		self.remove_button.setObjectName("remove_button")
		self.remove_button.setText("Remove column")
		self.remove_button.clicked.connect(page.remove_script)
		self.horizontal_layout.addWidget(self.remove_button)
		self.horizontal_layout.addItem(QtWidgets.QSpacerItem(1, 1, QtWidgets.QSizePolicy.Policy.Expanding, QtWidgets.QSizePolicy.Policy.Minimum))
		self.vertical_layout.addLayout(self.horizontal_layout)
		self.error_label = QtWidgets.QLabel(page)
		self.error_label.setObjectName("script_error")
		self.error_label.setText("")
		self.error_label.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter)
		self.vertical_layout.addWidget(self.error_label)


register_options_page(PluginOptionsPage)
update_columns()