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 

8from slixmpp import JID 

9from sqlalchemy.orm.exc import DetachedInstanceError 

10 

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

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

13from ...util.types import AnySession, Avatar 

14from .db import UpdateInfoMixin 

15 

16if TYPE_CHECKING: 

17 from ..pubsub import PepAvatar 

18 

19 

20class AvatarMixin(UpdateInfoMixin): 

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 jid: JID = NotImplemented 

29 session: AnySession = NotImplemented 

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 @property 

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

39 """ 

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

41 

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

43 later execution and is not blocking 

44 """ 

45 try: 

46 if self.stored.avatar is None: 

47 return None 

48 except DetachedInstanceError: 

49 self.merge() 

50 if self.stored.avatar is None: 

51 return None 

52 if self.stored.avatar.legacy_id is None: 

53 unique_id = None 

54 else: 

55 unique_id = self.stored.avatar.legacy_id 

56 return Avatar( 

57 unique_id=unique_id, 

58 url=self.stored.avatar.url, 

59 ) 

60 

61 @avatar.setter 

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

63 avatar = convert_avatar(avatar) 

64 if self._set_avatar_task: 

65 self._set_avatar_task.cancel() 

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

67 self._set_avatar_task = self.session.create_task(self.set_avatar(avatar)) 

68 

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

70 if self.avatar is None: 

71 return avatar is not None 

72 if avatar is None: 

73 return self.avatar is not None 

74 

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

76 return self.avatar.unique_id != avatar.unique_id 

77 

78 if ( 

79 self.avatar.url is not None 

80 and avatar.url is not None 

81 and self.avatar.url == avatar.url 

82 ): 

83 return await avatar_cache.url_modified(avatar.url) 

84 

85 if avatar.path is not None: 

86 cached = self.get_cached_avatar() 

87 if cached is not None: 

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

89 

90 return True 

91 

92 async def set_avatar( 

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

94 ) -> None: 

95 """ 

96 Set an avatar for this entity 

97 

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

99 ID 

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

101 it once used or not. 

102 """ 

103 avatar = convert_avatar(avatar) 

104 

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

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

107 avatar = Avatar( 

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

109 ) 

110 

111 try: 

112 if not await self.__has_changed(avatar): 

113 return 

114 except Exception: 

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

116 return 

117 

118 if avatar is None: 

119 cached_avatar = None 

120 else: 

121 try: 

122 cached_avatar = await avatar_cache.convert_or_get(avatar) 

123 except UnidentifiedImageError: 

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

125 cached_avatar = None 

126 except Exception: 

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

128 cached_avatar = None 

129 

130 if delete: 

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

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

133 else: 

134 avatar.path.unlink() 

135 

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

137 if not self._updating_info: 

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

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

140 orm.refresh(self.stored) 

141 

142 self.stored.avatar = stored_avatar 

143 self.commit() 

144 

145 self._post_avatar_update(cached_avatar) 

146 

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

148 try: 

149 if self.stored.avatar is None: 

150 return None 

151 except DetachedInstanceError: 

152 self.merge() 

153 if self.stored.avatar is None: 

154 return None 

155 return avatar_cache.get(self.stored.avatar) 

156 

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

158 cached_avatar = self.get_cached_avatar() 

159 if cached_avatar is None: 

160 return None 

161 from ..pubsub import PepAvatar 

162 

163 item = PepAvatar() 

164 item.set_avatar_from_cache(cached_avatar) 

165 return item 

166 

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

168 raise NotImplementedError 

169 

170 

171def convert_avatar( 

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

173) -> Avatar | None: 

174 if isinstance(avatar, Path): 

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

176 if isinstance(avatar, str): 

177 return Avatar(url=avatar) 

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

179 return None 

180 return avatar