Coverage for slidge/core/mixins/avatar.py: 88%

1import hashlib 

2from asyncio import Task 

3from logging import Logger 

4from pathlib import Path 

5from typing import TYPE_CHECKING, Optional 

6 

7import sqlalchemy as sa 

8import sqlalchemy.orm 

9from PIL import UnidentifiedImageError 

10 

11from ...db.avatar import AvatarDownloadError, CachedAvatar, avatar_cache 

12from ...db.models import Avatar as AvatarModel 

13from ...db.models import Contact, Room 

14from ...util.types import Avatar 

15from .base import SessionBound 

16from .db import UpdateInfoMixin 

17 

18if TYPE_CHECKING: 

19 from ..pubsub import PepAvatar 

20 

21 

22class AvatarMixin(UpdateInfoMixin, SessionBound): 

23 """ 

24 Mixin for XMPP entities that have avatars that represent them. 

25 

26 Both :py:class:`slidge.LegacyContact` and :py:class:`slidge.LegacyMUC` use 

27 :py:class:`.AvatarMixin`. 

28 """ 

29 

30 stored: Contact | Room 

31 log: Logger 

32 

33 def __init__(self) -> None: 

34 super().__init__() 

35 self._set_avatar_task: Task[None] | None = None 

36 

37 def __get_stored(self) -> AvatarModel | None: 

38 if self.stored.avatar_id is None: 

39 return None 

40 orm = sa.orm.object_session(self.stored) 

41 if ( 

42 orm is None 

43 and "avatar" in sa.orm.attributes.instance_state(self.stored).unloaded 

44 ): 

45 with self.xmpp.store.session() as orm, orm.no_autoflush: 

46 avatar = orm.execute( 

47 sa.select(AvatarModel).where( 

48 AvatarModel.id == self.stored.avatar_id 

49 ) 

50 ).scalar_one_or_none() 

51 sa.orm.attributes.set_committed_value(self.stored, "avatar", avatar) 

52 return self.stored.avatar 

53 

54 @property 

55 def avatar(self) -> Avatar | None: 

56 """ 

57 This property can be used to set or unset the avatar. 

58 

59 Unlike the awaitable :func:`.set_avatar`, it schedules the update for 

60 later execution and is not blocking 

61 """ 

62 stored = self.__get_stored() 

63 if stored is None: 

64 return None 

65 return Avatar(unique_id=stored.legacy_id, url=stored.url) 

66 

67 @avatar.setter 

68 def avatar(self, avatar: Avatar | Path | str | None) -> None: 

69 avatar = convert_avatar(avatar) 

70 if self._set_avatar_task: 

71 self._set_avatar_task.cancel() 

72 self.log.debug("Setting avatar with property") 

73 self._set_avatar_task = self.session.create_task( 

74 self.set_avatar(avatar), name=f"set avatar of {self}" 

75 ) 

76 

77 async def __has_changed(self, avatar: Avatar | None) -> bool: 

78 if self.avatar is None: 

79 return avatar is not None 

80 if avatar is None: 

81 return self.avatar is not None 

82 

83 if self.avatar.unique_id is not None and avatar.unique_id is not None: 

84 return self.avatar.unique_id != avatar.unique_id 

85 

86 if ( 

87 self.avatar.url is not None 

88 and avatar.url is not None 

89 and self.avatar.url == avatar.url 

90 ): 

91 return await avatar_cache.url_modified(avatar.url) 

92 

93 if avatar.path is not None: 

94 cached = self.get_cached_avatar() 

95 if cached is not None: 

96 return cached.path.read_bytes() != avatar.path.read_bytes() 

97 

98 return True 

99 

100 async def set_avatar( 

101 self, avatar: Avatar | Path | str | None = None, delete: bool = False 

102 ) -> None: 

103 """ 

104 Set an avatar for this entity 

105 

106 :param avatar: The avatar. Should ideally come with a legacy network-wide unique 

107 ID 

108 :param delete: If the avatar is provided as a Path, whether to delete 

109 it once used or not. 

110 """ 

111 avatar = convert_avatar(avatar) 

112 

113 if avatar is not None and avatar.unique_id is None and avatar.data is not None: 

114 self.log.debug("Hashing bytes to generate a unique ID") 

115 avatar = Avatar( 

116 data=avatar.data, unique_id=hashlib.sha512(avatar.data).hexdigest() 

117 ) 

118 

119 try: 

120 if not await self.__has_changed(avatar): 

121 return 

122 except AvatarDownloadError as e: 

123 self.log.warning("Could not determine if avatar has changed: %s", e) 

124 return 

125 except Exception: 

126 self.log.exception("Could not determine if avatar has changed, giving up") 

127 return 

128 

129 if avatar is None: 

130 cached_avatar = None 

131 else: 

132 try: 

133 cached_avatar = await avatar_cache.get(avatar, self.session) 

134 except UnidentifiedImageError: 

135 self.log.warning("%s is not a valid image", avatar) 

136 cached_avatar = None 

137 except AvatarDownloadError as e: 

138 self.log.warning("Could not fetch avatar %s: %s", avatar, e) 

139 cached_avatar = None 

140 except Exception: 

141 self.log.exception("Failed to set avatar '%s'", avatar) 

142 cached_avatar = None 

143 

144 if delete: 

145 if avatar is None or avatar.path is None: 

146 self.log.warning("Requested avatar path delete, but no path provided") 

147 else: 

148 avatar.path.unlink() 

149 

150 stored_avatar = None if cached_avatar is None else cached_avatar.stored 

151 if not self._updating_info: 

152 with self.xmpp.store.session() as orm, orm.no_autoflush: 

153 self.stored = orm.merge(self.stored) 

154 orm.refresh(self.stored) 

155 

156 self.stored.avatar = stored_avatar 

157 self.commit() 

158 

159 self._post_avatar_update(cached_avatar) 

160 

161 def get_cached_avatar(self) -> Optional["CachedAvatar"]: 

162 stored = self.__get_stored() 

163 if stored is None: 

164 return None 

165 return avatar_cache.from_stored(stored) 

166 

167 def get_avatar(self) -> Optional["PepAvatar"]: 

168 cached_avatar = self.get_cached_avatar() 

169 if cached_avatar is None: 

170 return None 

171 from ..pubsub import PepAvatar 

172 

173 item = PepAvatar() 

174 item.set_avatar_from_cache(cached_avatar) 

175 return item 

176 

177 def _post_avatar_update(self, cached_avatar: Optional["CachedAvatar"]) -> None: 

178 raise NotImplementedError 

179 

180 

181def convert_avatar( 

182 avatar: Avatar | Path | str | None, unique_id: str | None = None 

183) -> Avatar | None: 

184 if isinstance(avatar, Path): 

185 return Avatar(path=avatar, unique_id=unique_id) 

186 if isinstance(avatar, str): 

187 return Avatar(url=avatar) 

188 if avatar is None or all(x is None for x in avatar): 

189 return None 

190 return avatar