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

1import hashlib 

2from asyncio import Task 

3from logging import Logger 

4from pathlib import Path 

5from typing import TYPE_CHECKING, Optional 

6 

7from PIL import UnidentifiedImageError 

8 

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

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

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

12from ...util.types import Avatar 

13from .base import SessionBound 

14from .db import UpdateInfoMixin 

15 

16if TYPE_CHECKING: 

17 from ..pubsub import PepAvatar 

18 

19 

20class AvatarMixin(UpdateInfoMixin, SessionBound): 

21 """ 

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

23 

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

25 :py:class:`.AvatarMixin`. 

26 """ 

27 

28 stored: Contact | Room 

29 log: Logger 

30 

31 def __init__(self) -> None: 

32 super().__init__() 

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

34 

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

36 self.stored.ensure_relation("avatar", self.xmpp.store.session) 

37 return self.stored.avatar 

38 

39 @property 

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

41 """ 

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

43 

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

45 later execution and is not blocking 

46 """ 

47 stored = self.__get_stored() 

48 if stored is None: 

49 return None 

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

51 

52 @avatar.setter 

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

54 avatar = convert_avatar(avatar) 

55 if self._set_avatar_task: 

56 self._set_avatar_task.cancel() 

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

58 self._set_avatar_task = self.session.create_task( 

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

60 ) 

61 

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

63 if self.avatar is None: 

64 return avatar is not None 

65 if avatar is None: 

66 return self.avatar is not None 

67 

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

69 return self.avatar.unique_id != avatar.unique_id 

70 

71 if ( 

72 self.avatar.url is not None 

73 and avatar.url is not None 

74 and self.avatar.url == avatar.url 

75 ): 

76 return await avatar_cache.url_modified(avatar.url) 

77 

78 if avatar.path is not None: 

79 cached = self.get_cached_avatar() 

80 if cached is not None: 

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

82 

83 return True 

84 

85 async def set_avatar( 

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

87 ) -> None: 

88 """ 

89 Set an avatar for this entity 

90 

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

92 ID 

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

94 it once used or not. 

95 """ 

96 avatar = convert_avatar(avatar) 

97 

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

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

100 avatar = Avatar( 

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

102 ) 

103 

104 try: 

105 if not await self.__has_changed(avatar): 

106 return 

107 except AvatarDownloadError as e: 

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

109 return 

110 except Exception: 

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

112 return 

113 

114 if avatar is None: 

115 cached_avatar = None 

116 else: 

117 try: 

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

119 except UnidentifiedImageError: 

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

121 cached_avatar = None 

122 except AvatarDownloadError as e: 

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

124 cached_avatar = None 

125 except Exception: 

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

127 cached_avatar = None 

128 

129 if delete: 

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

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

132 else: 

133 avatar.path.unlink() 

134 

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

136 if not self._updating_info: 

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

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

139 orm.refresh(self.stored) 

140 

141 self.stored.avatar = stored_avatar 

142 self.commit() 

143 

144 self._post_avatar_update(cached_avatar) 

145 

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

147 stored = self.__get_stored() 

148 if stored is None: 

149 return None 

150 return avatar_cache.from_stored(stored) 

151 

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

153 cached_avatar = self.get_cached_avatar() 

154 if cached_avatar is None: 

155 return None 

156 from ..pubsub import PepAvatar 

157 

158 item = PepAvatar() 

159 item.set_avatar_from_cache(cached_avatar) 

160 return item 

161 

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

163 raise NotImplementedError 

164 

165 

166def convert_avatar( 

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

168) -> Avatar | None: 

169 if isinstance(avatar, Path): 

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

171 if isinstance(avatar, str): 

172 return Avatar(url=avatar) 

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

174 return None 

175 return avatar