Coverage for shrinky/__init__.py: 47%

1""" shrinky module """ 

2 

3import os 

4from pathlib import Path 

5from typing import Dict, Optional, Tuple, Union 

6 

7from loguru import logger 

8from PIL import Image 

9 

10 

11DEFAULT_GEOMETRY = 2000 

12 

13def parse_geometry(geometry_input: str) -> Tuple[Optional[int], Optional[int]]: 

14 """parse the geometry provided""" 

15 if "x" in geometry_input: 

16 x_value, y_value = geometry_input.split("x") 

17 if x_value == "": 

18 x_result = None 

19 else: 

20 x_result = int(x_value) 

21 

22 if y_value == "": 

23 y_result = None 

24 else: 

25 y_result = int(y_value) 

26 return (x_result, y_result) 

27 return (0, 0) 

28 

29 

30def set_geometry(geometry_value: Optional[str]) -> Tuple[int, int]: 

31 """geometry handler""" 

32 

33 if geometry_value is None: 

34 return (DEFAULT_GEOMETRY, DEFAULT_GEOMETRY) 

35 

36 max_x, max_y = parse_geometry(geometry_value) 

37 if max_x is None: 

38 max_x = DEFAULT_GEOMETRY 

39 if max_y is None: 

40 max_y = DEFAULT_GEOMETRY 

41 logger.debug("Setting geometry to {}x{}", max_x, max_y) 

42 return (max_x, max_y) 

43 

44def new_filename(original_filename: Path, output_type: Optional[str]) -> Path: 

45 """generates a new filename based on the path""" 

46 logger.debug(f"{original_filename=}") 

47 

48 basename = ".".join(original_filename.resolve().name.split(".")[:-1]) 

49 if output_type is not None: 

50 logger.debug("Setting output type to {}", output_type.lower()) 

51 newname = f"{basename}.{output_type.lower()}" 

52 else: 

53 extension = get_extension(original_filename) 

54 newname = f"{basename}-shrink.{extension}" 

55 

56 return Path(f"{original_filename.parent}/{newname}").resolve() 

57 

58 

59def get_extension(filename: Path) -> str: 

60 """gets the file extension is a hacky way""" 

61 if "." not in filename.name: 

62 raise ValueError("Can't have an extension when there's no dot!") 

63 return filename.resolve().name.split(".")[-1] 

64 

65 

66class ShrinkyImage: 

67 """does all the things""" 

68 

69 def __init__(self, source_path: Path): 

70 """loads the image""" 

71 self.source_path = source_path 

72 self.image = Image.open(source_path.open("rb")) 

73 logger.debug("Dims: {}x{}", self.image.width, self.image.height) 

74 logger.info( 

75 "Original file size: {}", os.stat(self.source_path.resolve()).st_size 

76 ) 

77 

78 def resize_image( 

79 self, 

80 new_width: int, 

81 new_height: int, 

82 source_image: Optional[Image.Image] = None, 

83 ) -> Image.Image: 

84 """resizes an image, doesn't modify the source image""" 

85 if source_image is None: 

86 source_image = self.image 

87 

88 tmpimage = source_image.copy() 

89 

90 if source_image.width > new_width or source_image.height > new_height: 

91 logger.debug( 

92 "Thumbnailing from {}x{} to {}x{}", 

93 source_image.width, 

94 source_image.height, 

95 new_width, 

96 new_height, 

97 ) 

98 tmpimage.thumbnail((new_width, new_height)) 

99 return tmpimage 

100 

101 def write_image( 

102 self, output_filename: Path, force_overwrite: bool = False, quality: int = -1 

103 ) -> bool: 

104 """writes the file to disk""" 

105 if output_filename.exists() and not force_overwrite: 

106 logger.error("{} already exists, bailing", output_filename.resolve()) 

107 return False 

108 

109 args: Dict[str, Union[str, int]] = {} 

110 

111 if get_extension(output_filename).lower() in ("jpg", "jpeg"): 

112 # set jpeg quality 

113 if quality is not None and quality >= 0: 

114 args["quality"] = quality 

115 

116 if self.image.mode != "RGB": 

117 self.image = self.image.convert("RGB") 

118 

119 with output_filename.open("wb") as output_image: 

120 logger.info("Writing {}", output_image.name) 

121 self.image.save(output_image, **args) # type: ignore 

122 

123 new_size = os.stat(output_filename.resolve()).st_size 

124 logger.info("New size: {}", new_size) 

125 return True