Coverage for slidge/core/mixins/avatar.py: 89%
119 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-10 12:26 +0000
1import asyncio
2import hashlib
3from logging import Logger
4from pathlib import Path
5from typing import TYPE_CHECKING, Optional
7import sqlalchemy as sa
8from PIL import UnidentifiedImageError
10from ...db.avatar import AvatarDownloadError, CachedAvatar, avatar_cache
11from ...db.models import Avatar as AvatarModel
12from ...db.models import Contact, Room, Space
13from ...util.types import Avatar, ExtraAttrMapping
14from .db import UpdateInfoMixin
16if TYPE_CHECKING:
17 from ..pubsub import PepAvatar
20class ImageMixin[T: ExtraAttrMapping, PropsT](UpdateInfoMixin[T, PropsT]):
21 """
22 Mixin for XMPP entities that have images (avatar or banner) that represent them.
24 Both :py:class:`slidge.LegacyContact` and :py:class:`slidge.LegacyRoom` use
25 :py:class:`.AvatarMixin`.
26 """
28 _stored: Contact | Room | Space
29 log: Logger
31 def _get_stored_image(self, attr: str) -> AvatarModel | None:
32 self._stored.ensure_relation(attr, self.xmpp.store.session)
33 return getattr(self._stored, attr) # type:ignore[no-any-return]
35 @property
36 def avatar(self) -> Avatar | None:
37 """
38 This property can be used to set or unset the avatar.
39 """
40 stored = self._get_stored_image("avatar")
41 if stored is None:
42 return None
43 return stored.to_meta()
45 @avatar.setter
46 def avatar(self, avatar: Avatar | Path | str | None) -> None:
47 self._image_setter(avatar, "avatar")
49 def _image_setter(
50 self, avatar: Avatar | Path | str | None, attr: str = "avatar"
51 ) -> None:
52 new = convert_avatar(avatar)
53 old = self._get_stored_image(attr)
55 if new is None:
56 if old is None:
57 self._session.log.debug(f"early cancel {attr} None")
58 return
59 else:
60 if old is not None and new.unique_id and old.legacy_id == new.unique_id:
61 self._session.log.debug(f"early cancel {attr}")
62 return
64 self._session.create_task(
65 self._set_image(avatar, attr=attr),
66 name=self._set_image_task_name(attr),
67 )
69 def _set_image_task_name(self, attr: str = "avatar") -> str:
70 return f"set {attr} of {self._stored.__class__} #{self._stored.id}"
72 @property
73 def _set_avatar_task(self) -> asyncio.Task[None] | None:
74 return self._session.get_task(self._set_image_task_name("avatar"))
76 async def __detect_changes(self, old: Avatar | None, avatar: Avatar | None) -> bool:
77 if old is None:
78 return avatar is not None
79 if avatar is None:
80 return old is not None
82 if old.unique_id is not None and avatar.unique_id is not None:
83 return old.unique_id != avatar.unique_id
85 if old.url is not None and avatar.url is not None and old.url == avatar.url:
86 return await avatar_cache.url_modified(avatar.url)
88 if avatar.path is not None:
89 cached = self.get_cached_avatar()
90 if cached is not None:
91 return cached.path.read_bytes() != avatar.path.read_bytes()
93 return True
95 async def _set_image(
96 self,
97 avatar: Avatar | Path | str | None = None,
98 *,
99 attr: str = "avatar",
100 ) -> None:
101 """
102 Set an avatar for this entity
104 :param avatar: The avatar. Should ideally come with a legacy network-wide unique
105 ID
106 :param delete: If the avatar is provided as a Path, whether to delete
107 it once used or not.
108 """
109 old = self._get_stored_image(attr)
110 new = convert_avatar(avatar)
111 await self._process_image(
112 None if old is None else old.to_meta(), new, attr=attr
113 )
115 async def _process_image(
116 self, old: Avatar | None, new: Avatar | None, *, attr: str
117 ) -> None:
118 if new is not None and new.unique_id is None and new.data is not None:
119 self.log.debug("Hashing bytes to generate a unique ID")
120 new = Avatar(data=new.data, unique_id=hashlib.sha512(new.data).hexdigest())
122 try:
123 if not await self.__detect_changes(old, new):
124 return
125 except AvatarDownloadError as e:
126 self.log.warning("Could not determine if avatar has changed: %s", e)
127 return
128 except Exception:
129 self.log.exception("Could not determine if avatar has changed, giving up")
130 return
132 if new is None:
133 cached_avatar = None
134 else:
135 try:
136 cached_avatar = await avatar_cache.get(new, self._session)
137 except UnidentifiedImageError:
138 self.log.warning("%s is not a valid image", new)
139 cached_avatar = None
140 except AvatarDownloadError as e:
141 self.log.warning("Could not fetch avatar %s: %s", new, e)
142 cached_avatar = None
143 except Exception:
144 self.log.exception("Failed to set avatar '%s'", new)
145 cached_avatar = None
147 if cached_avatar is None:
148 stored = pk = None
149 else:
150 stored = cached_avatar._stored
151 pk = stored.id
153 with self.xmpp.store.session() as orm:
154 kwargs = {f"{attr}_id": pk}
155 cls = self._stored.__class__
156 orm.execute(
157 sa.update(cls).where(cls.id == self._stored.id).values(**kwargs)
158 )
159 orm.commit()
160 self._stored._set_committed_value(f"{attr}", stored)
161 self._stored._set_committed_value(f"{attr}_id", pk)
163 self._post_image_update(cached_avatar, attr=attr)
165 def get_cached_avatar(self, name: str = "avatar") -> Optional["CachedAvatar"]:
166 stored = self._get_stored_image(name)
167 if stored is None:
168 return None
169 return avatar_cache.from_stored(stored)
171 def get_avatar(self) -> Optional["PepAvatar"]:
172 cached_avatar = self.get_cached_avatar()
173 if cached_avatar is None:
174 return None
175 from ..pubsub import PepAvatar
177 item = PepAvatar()
178 item.set_avatar_from_cache(cached_avatar)
179 return item
181 def _post_image_update(
182 self, cached_avatar: Optional["CachedAvatar"], *, attr: str = "avatar"
183 ) -> None:
184 raise NotImplementedError
187def convert_avatar(
188 avatar: Avatar | Path | str | None, unique_id: str | None = None
189) -> Avatar | None:
190 if isinstance(avatar, Path):
191 return Avatar(path=avatar, unique_id=unique_id)
192 if isinstance(avatar, str):
193 return Avatar(url=avatar)
194 if avatar is None or all(x is None for x in avatar):
195 return None
196 return avatar