Coverage for slidge / util / lottie.py: 0%

1""" 

2This module implements vector animated stickers in the lottie format to webp images. 

3 

4 

5""" 

6 

7import asyncio 

8import logging 

9import warnings 

10from pathlib import Path 

11 

12import aiohttp 

13 

14try: 

15 from rlottie_python.rlottie_wrapper import LottieAnimation 

16except ImportError: 

17 LottieAnimation = None # type: ignore[assignment,misc] 

18 

19from ..core import config 

20from .types import LegacyAttachment 

21 

22 

23async def from_url( 

24 url: str, sticker_id: str, http: aiohttp.ClientSession 

25) -> LegacyAttachment: 

26 """ 

27 Get a webp attachment from a URL. 

28 

29 :param url: URL where the lottie sticker can be downloaded. 

30 :param sticker_id: A unique identifier for this sticker. 

31 :param http: The aiohttp.ClientSession used to download the sticker. 

32 

33 :return: A `LegacyAttachment` with the sticker in the webp format if 

34 `config.CONVERT_STICKERS == True`, in the original lottie format 

35 otherwise. 

36 """ 

37 if not config.CONVERT_STICKERS: 

38 return _attachment(sticker_id, url=url) 

39 lottie_path = sticker_path(sticker_id).with_suffix(".json") 

40 await _download(lottie_path, url, http) 

41 webp_path = sticker_path(sticker_id).with_suffix(".webp") 

42 return await convert(lottie_path, webp_path, sticker_id) 

43 

44 

45async def from_path(path: Path, sticker_id: str) -> LegacyAttachment: 

46 """ 

47 Get a webp attachment from a path. 

48 

49 :param path: path to the lottie sticker file. 

50 :param sticker_id: A unique identifier for this sticker. 

51 

52 :return: A `LegacyAttachment` with the sticker in the webp format if 

53 `config.CONVERT_STICKERS == True`, in the original lottie format 

54 otherwise. 

55 """ 

56 if not config.CONVERT_STICKERS: 

57 return _attachment(sticker_id, path=path) 

58 out_path = sticker_path(sticker_id).with_suffix(".webp") 

59 return await convert(path, out_path, sticker_id) 

60 

61 

62def sticker_path(sticker_id: str = "") -> Path: 

63 """ 

64 Get the path where a sticker is meant be stored 

65 

66 :param sticker_id: A unique ID for this sticker 

67 :return: A `path`. It has no suffix, so you might want to use `Path.with_suffix()` 

68 before writing into it. 

69 """ 

70 root = config.HOME_DIR / "lottie" 

71 root.mkdir(exist_ok=True) 

72 return root / sticker_id 

73 

74 

75def _attachment( 

76 sticker_id: str, path: Path | None = None, url: str | None = None 

77) -> LegacyAttachment: 

78 content_type = "image/webp" if path is not None and path.suffix == ".webp" else None 

79 return LegacyAttachment( 

80 url=url, 

81 path=path, 

82 legacy_file_id="lottie-" + sticker_id, 

83 disposition="inline", 

84 content_type=content_type, 

85 is_sticker=True, 

86 ) 

87 

88 

89async def _download(path: Path, url: str, http: aiohttp.ClientSession) -> None: 

90 async with _sticker_download_lock: 

91 if path.exists(): 

92 return 

93 with path.open("wb") as fp: 

94 try: 

95 resp = await http.get(url, chunked=True) 

96 async for chunk in resp.content: 

97 fp.write(chunk) 

98 except ValueError as e: 

99 if e.args[0] != "Chunk too big": 

100 raise 

101 # not sure why this happens but it does sometimes 

102 resp = await http.get(url) 

103 fp.write(await resp.content.read()) 

104 

105 

106async def convert( 

107 input_path: Path, 

108 output_path: Path, 

109 sticker_id: str, 

110 width: int = 256, 

111 height: int = 256, 

112) -> LegacyAttachment: 

113 async with _sticker_conversion_lock: 

114 if not output_path.exists(): 

115 if LottieAnimation is None: 

116 warnings.warn( 

117 "Cannot convert stickers, rlottie-python is not available." 

118 ) 

119 return _attachment(sticker_id, path=input_path) 

120 

121 log.debug("Converting sticker %s to video", output_path.stem) 

122 if input_path.suffix == ".json": 

123 animation = LottieAnimation.from_file(str(input_path)) 

124 else: 

125 animation = LottieAnimation.from_tgs(str(input_path)) 

126 animation.save_animation(str(output_path), width=width, height=height) 

127 

128 return _attachment(sticker_id, path=output_path) 

129 

130 

131async def _main() -> None: 

132 # small entrypoint to easily test that it works 

133 import sys 

134 

135 source, destination = sys.argv[1:] 

136 config.CONVERT_STICKERS = True 

137 await convert(Path(source), Path(destination), "") 

138 

139 

140_sticker_conversion_lock: asyncio.Lock = asyncio.Lock() 

141_sticker_download_lock: asyncio.Lock = asyncio.Lock() 

142 

143 

144log = logging.getLogger(__name__) 

145 

146 

147__all__ = ("from_url",) 

148 

149if __name__ == "__main__": 

150 asyncio.run(_main())