Adding FileGroup.set_tags which replaces several tag manipulation methods for FileGroups.

wrolpi/files/api.py

@@ -305,11 +305,6 @@ async def post_untag_file_group(_, body: schema.TagFileGroupPost):
     return response.empty(HTTPStatus.NO_CONTENT)
 
 
-# {
-#   '/media/directory/sub-dir/the-file-name.suffix': 2,  # The chunk number we will receive next.
-# }
-
-
 @files_bp.post('/upload')
 async def post_upload(request: Request):
     """Accepts a multipart/form-data request to upload a single file.

wrolpi/files/models.py

@@ -3,6 +3,7 @@ import pathlib
 import shutil
 import urllib.parse
 from datetime import datetime
+from functools import singledispatchmethod
 from typing import List, Type, Optional, Iterable
 
 from sqlalchemy import Column, String, Computed, BigInteger, Boolean
@@ -17,7 +18,7 @@ from wrolpi.downloader import Download
 from wrolpi.errors import FileGroupIsTagged, UnknownFile
 from wrolpi.files import indexers
 from wrolpi.media_path import MediaPathType
-from wrolpi.tags import Tag, TagFile
+from wrolpi.tags import Tag, TagFile, save_tags_config, sync_tags_directory
 from wrolpi.vars import PYTEST
 
 logger = logger.getChild(__name__)
@@ -144,12 +145,56 @@ class FileGroup(ModelHelper, Base):
         self.viewed = now()
 
     @optional_session
-    def add_tag(self, tag_id_or_name: int | str, session: Session = None) -> TagFile:
-        return Tag.tag_file_group(self.id, tag_id_or_name, session)
+    def set_tags(self, tag_names_or_ids: Iterable[str | int], session: Session = None):
+        """Insert or Delete TagFiles as necessary to match the provided tags for this FileGroup."""
+        tag_names_or_ids = list(tag_names_or_ids)
+        if tag_names_or_ids and isinstance(tag_names_or_ids[0], str):
+            tag_ids = {Tag.get_id_by_name(i) for i in tag_names_or_ids}
+        else:
+            tag_ids = set(tag_names_or_ids)
 
+        existing_tag_ids = {i.tag_id for i in self.tag_files}
+        if new_tag_ids := (tag_ids - existing_tag_ids):
+            for tag_id in new_tag_ids:
+                tag_file = TagFile(file_group_id=self.id, tag_id=tag_id)
+                self.tag_files.append(tag_file)
+            session.flush(self.tag_files)
+
+        if deleted_tags := (existing_tag_ids - tag_ids):
+            session.query(TagFile).filter(
+                TagFile.file_group_id == self.id,
+                TagFile.tag_id.in_(deleted_tags),
+            ).delete(synchronize_session=False)
+
+        # Save changes to config.
+        save_tags_config.activate_switch()
+        sync_tags_directory.activate_switch()
+
+        return self.tag_files
+
+    @singledispatchmethod
     @optional_session
-    def untag(self, tag_id_or_name: int | str, session: Session = None):
-        Tag.untag_file_group(self.id, tag_id_or_name, session)
+    def add_tag(self, tag_id: int | str, session: Session = None) -> TagFile:
+        tags = [i.tag_id for i in self.tag_files]
+        tags.append(tag_id)
+        tag_files = self.set_tags(tags, session=session)
+        return next(i for i in tag_files if i.tag_id == tag_id)
+
+    @add_tag.register
+    def _(self, tag_name: str, session: Session = None) -> TagFile:
+        tag_id = Tag.get_id_by_name(tag_name, session=session)
+        return self.add_tag(tag_id, session=session)
+
+    @singledispatchmethod
+    @optional_session
+    def untag(self, tag_id: int | str, session: Session = None):
+        tags = [i.tag_id for i in self.tag_files if i.tag_id != tag_id]
+        self.set_tags(tags, session=session)
+
+    @untag.register
+    def _(self, tag_name: str, session: Session = None) -> TagFile:
+        tag_id = Tag.get_id_by_name(tag_name, session=session)
+        return self.untag(tag_id, session=session)
 
     def append_files(self, *paths: pathlib.Path):
         """Add all `paths` to this FileGroup.files."""

wrolpi/files/test/test_lib.py

@@ -16,6 +16,7 @@ from PIL import Image
 import wrolpi.common
 from modules.videos import Video
 from wrolpi.common import timer
+from wrolpi.conftest import await_switches
 from wrolpi.dates import now
 from wrolpi.errors import InvalidFile, UnknownDirectory, FileGroupIsTagged, NoPrimaryFile
 from wrolpi.files import lib, indexers

wrolpi/files/test/test_models.py

@@ -1,5 +1,8 @@
 import pytest
 
+from wrolpi.files import lib
+from wrolpi.files.models import FileGroup
+
 
 @pytest.mark.asyncio
 async def test_file_group_move(async_client, test_session, test_directory, video_factory):
@@ -19,3 +22,23 @@ async def test_file_group_move(async_client, test_session, test_directory, video
         (test_directory / 'videos/NO CHANNEL/new name.jpg'),
         (test_directory / 'videos/NO CHANNEL/new name.mp4'),
     ]
+
+
+@pytest.mark.asyncio
+async def test_add_tag_overload(async_client, test_session, make_files_structure, tag_factory):
+    """FileGroup.add_tag and FileGroup.untag use `singeldispatch`."""
+    make_files_structure({'foo.txt': 'foo contents'})
+    await lib.refresh_files()
+    foo: FileGroup = test_session.query(FileGroup).one()
+
+    one, two = await tag_factory(), await tag_factory()
+
+    foo.add_tag(one.name)
+    assert set(foo.tag_names) == {'one'}, 'Tag should have been found by name'
+    foo.add_tag(two.id)
+    assert set(foo.tag_names) == {'one', 'two'}, 'Tag should have been found by ID.'
+
+    foo.untag(one.name)
+    assert set(foo.tag_names) == {'two'}, 'Tag should have been found by name'
+    foo.untag(two.id)
+    assert not foo.tag_names, 'Tag should have been found by ID.'

wrolpi/tags.py

@@ -4,7 +4,9 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from typing import List, Dict, Tuple, Optional
 
-from sqlalchemy import Column, Integer, String, ForeignKey, BigInteger
+import cachetools
+from cachetools.keys import hashkey
+from sqlalchemy import Column, Integer, String, ForeignKey, BigInteger, event
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import relationship, Session
 
@@ -14,7 +16,7 @@ from wrolpi.common import ModelHelper, Base, logger, ConfigFile, get_media_direc
 from wrolpi.dates import TZDateTime
 from wrolpi.db import optional_session, get_db_curs, get_db_session
 from wrolpi.downloader import save_downloads_config
-from wrolpi.errors import UnknownTag, UsedTag, InvalidTag, FileGroupAlreadyTagged, RefreshConflict, NoPrimaryFile
+from wrolpi.errors import UnknownTag, UsedTag, InvalidTag, RefreshConflict, NoPrimaryFile
 from wrolpi.events import Events
 from wrolpi.switches import register_switch_handler, ActivateSwitchMethod
 from wrolpi.vars import PYTEST
@@ -51,6 +53,22 @@ class TagFile(ModelHelper, Base):
             .one_or_none()
 
 
+# Tag.get_id_by_name
+def get_id_by_name_key(klass: Base, name: str, session: Session = None):
+    return hashkey(name)
+
+
+get_id_by_name_cache = cachetools.LRUCache(maxsize=1_000)
+
+
+# Tag.get_name_by_id
+def get_name_by_id_key(klass: Base, id_: int, session: Session = None):
+    return hashkey(id_)
+
+
+get_name_by_id_cache = cachetools.LRUCache(maxsize=1_000)
+
+
 class Tag(ModelHelper, Base):
     __tablename__ = 'tag'
     id = Column(Integer, primary_key=True)
@@ -74,55 +92,32 @@ class Tag(ModelHelper, Base):
             color=self.color,
         )
 
-    @staticmethod
+    @classmethod
+    @cachetools.cached(cache=get_id_by_name_cache, key=get_id_by_name_key)
     @optional_session
-    def get_tag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None) -> Optional[TagFile]:
-        if isinstance(tag_id_or_name, int):
-            return TagFile.get_by_primary_keys(file_group_id, tag_id_or_name, session)
+    def get_id_by_name(cls, name: str, session: Session = None) -> int:
+        """Returns a Tag's ID, if a Tag matches the provided name.
 
-        return TagFile.get_by_tag_name(file_group_id, tag_id_or_name, session)
+        @raise UnknownTag: If no Tag matches the provided name."""
+        if tag := session.query(Tag).filter_by(name=name).one_or_none():
+            return tag.id
+        raise UnknownTag(f'No Tag with name={name}')
 
-    @staticmethod
+    @classmethod
+    @cachetools.cached(cache=get_name_by_id_cache, key=get_name_by_id_key)
     @optional_session
-    def tag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None) -> TagFile:
-        """Add a TagFile for the provided FileGroup and this Tag.
+    def get_name_by_id(cls, id_: int, session: Session = None) -> str:
+        """Returns a Tag's name, if a Tag matches the provided ID.
 
-        @warning: Commits the session to keep the config in sync."""
-        existing = Tag.get_tag_file_group(file_group_id, tag_id_or_name, session)
+        @raise UnknownTag: If no Tag matches the provided ID."""
+        if tag := session.query(Tag).filter_by(id=id_).one_or_none():
+            return tag.id
+        raise UnknownTag(f'No Tag with id={id_}')
 
-        if existing:
-            raise FileGroupAlreadyTagged('Tag already used')
-
-        tag = Tag.get_by_name(tag_id_or_name, session) if isinstance(tag_id_or_name, str) \
-            else Tag.get_by_id(tag_id_or_name, session)
-        tag_file = TagFile(file_group_id=file_group_id, tag_id=tag.id)
-        session.add(tag_file)
-        session.flush([tag_file])
-        session.commit()
-        logger.info(f'Tagged FileGroup {file_group_id} with Tag {tag_id_or_name}')
-
-        # Save changes to config.
-        save_tags_config.activate_switch()
-        sync_tags_directory.activate_switch()
-        return tag_file
-
-    @staticmethod
-    @optional_session
-    def untag_file_group(file_group_id: int, tag_id_or_name: int | str, session: Session = None):
-        """Remove the record of a Tag applied to the FileGroup.
-
-        @warning: Commits the session to keep config in sync."""
-        tag_file = Tag.get_tag_file_group(file_group_id, tag_id_or_name, session)
-        if tag_file:
-            session.delete(tag_file)
-            session.commit()
-            logger.debug(f'Deleted TagFile: {tag_file}')
-
-            # Save changes to config.
-            save_tags_config.activate_switch()
-            sync_tags_directory.activate_switch()
-        else:
-            logger.warning(f'Could not find TagFile for FileGroup.id={file_group_id}/Tag.id={tag_id_or_name}')
+    @classmethod
+    def invalidate_cache(cls):
+        get_id_by_name_cache.clear()
+        get_name_by_id_cache.clear()
 
     @staticmethod
     @optional_session
@@ -136,7 +131,7 @@ class Tag(ModelHelper, Base):
         if tag := Tag.get_by_name(name, session):
             return tag
 
-        raise UnknownTag(f'Unknown Tag with name={name}')
+        raise UnknownTag(f'No Tag with name={name}')
 
     @staticmethod
     @optional_session
@@ -150,7 +145,7 @@ class Tag(ModelHelper, Base):
         if tag := Tag.get_by_id(id_, session):
             return tag
 
-        raise UnknownTag(f'Unknown Tag with id={id_}')
+        raise UnknownTag(f'No Tag with id={id_}')
 
     def has_relations(self) -> bool:
         """Returns True if this Tag has been used with any FileGroups or Zim Entries."""
@@ -254,6 +249,14 @@ class Tag(ModelHelper, Base):
         sync_tags_directory.activate_switch()
 
 
+@event.listens_for(Tag, 'after_insert')
+@event.listens_for(Tag, 'after_update')
+@event.listens_for(Tag, 'after_delete')
+def invalidate_cache(mapper, connection, target):
+    """Clear Tag cache when Tags are changed."""
+    Tag.invalidate_cache()
+
+
 @dataclass
 class TagsConfigValidator:
     tag_files: list[list] = field(default_factory=list)

wrolpi/test/test_tags.py

@@ -6,7 +6,7 @@ import yaml
 
 from wrolpi import tags
 from wrolpi.common import is_hardlinked, walk
-from wrolpi.errors import FileGroupIsTagged, InvalidTag, UnknownTag, FileGroupAlreadyTagged, UsedTag
+from wrolpi.errors import FileGroupIsTagged, InvalidTag, UnknownTag, UsedTag
 from wrolpi.files import lib as files_lib
 from wrolpi.files.models import FileGroup
 from wrolpi.tags import TagFile, Tag
@@ -20,19 +20,18 @@ async def test_tags_file_group_json(async_client, test_session, make_files_struc
     await files_lib.refresh_files()
     file_group: FileGroup = test_session.query(FileGroup).one()
 
-    Tag.tag_file_group(file_group.id, tag_one.id)
+    file_group.add_tag(tag_one.id)
     assert file_group.__json__()['tags'] == ['one']
 
-    file_group.add_tag(tag_two.id)
+    file_group.add_tag(tag_two.name)
     assert file_group.__json__()['tags'] == ['one', 'two']
 
-    with pytest.raises(FileGroupAlreadyTagged):
-        # Cannot tag a FileGroup with the same Tag twice.
-        file_group.add_tag(tag_two.id)
-
     file_group.untag(tag_one.id)
     assert file_group.__json__()['tags'] == ['two']
 
+    file_group.untag(tag_two.name)
+    assert file_group.__json__()['tags'] == []
+
 
 @pytest.mark.asyncio
 async def test_tags_model(async_client, test_session, make_files_structure, tag_factory, example_pdf):
@@ -66,13 +65,13 @@ async def test_tags_file_group(async_client, test_session, make_files_structure,
     tag_one = await tag_factory()
     assert len(tag_one.tag_files) == 0
 
-    Tag.tag_file_group(video_group.id, tag_one.id)
+    video_group.add_tag(tag_one.id)
     test_session.commit()
     assert len(tag_one.tag_files) == 1
     assert sorted([i.tag.name for i in video_group.tag_files]) == ['one', ]
 
     tag_two = await tag_factory()
-    Tag.tag_file_group(video_group.id, tag_two.id)
+    video_group.add_tag(tag_two.id)
     test_session.commit()
     assert len(tag_one.tag_files) == 1
     assert len(tag_two.tag_files) == 1