formatitko/images.py

from typing import List
import os
import shutil
import subprocess
from PIL import Image

class ImageProcessor:
	def __init__(self, public_dir: str, web_path: str, *lookup_dirs: List[str]):
		self.public_dir = public_dir
		self.lookup_dirs = lookup_dirs
		self.web_path = web_path if web_path[-1] != "/" else web_path[:-1]
		if not os.path.exists(self.public_dir):
			os.mkdir(self.public_dir)

	def process_image(self, input_filename: str, format: str, source_dir: str, relative: bool=True, width: int=None, height:int=None, quality: int=None, dpi: int=None, fit: bool=True) -> str:
		name = os.path.basename(input_filename)
		base, ext = os.path.splitext(name)
		ext = ext[1:]
		full_path = self.find_image(input_filename, [source_dir])
		if full_path is None:
			raise FileNotFoundError(f'Image {input_filename} not found.')

		# Generate filename from arguments
		suffix = ""
		geometry = None
		if width is not None or height is not None:
			geometry = f'{width if width is not None else ""}x{height if height is not None else ""}{"" if fit else "!"}'
			suffix += "_"+geometry
		if quality is not None:
			suffix += f'_q{quality}'
		target_name = base+suffix+"."+format
		target_path = self.public_dir + "/" + target_name
	
		# Only regenerate if the file doesn't already exist.
		if not os.path.isfile(target_path):

			# If the format is the same or it is just a different extension for
			# the same format, just copy it.
			if (((ext == format)
		     or (ext == "epdf" and format == "pdf")
			 or (ext == "jpg" and format == "jpeg"))
			 and width is None and height is None and quality is None and dpi is None):
				shutil.copyfile(full_path, target_path)

			# Try to find the converted filename in lookup_dirs, if you find
			# it, don't convert, just copy.
			elif self.find_image(target_name, [source_dir]):
				shutil.copyfile(self.find_image(target_name, [source_dir]), target_path)

			# Convert SVGs using inkscape
			elif ext == "svg":
				width_arg = ['--export-width', str(width)] if width is not None else []
				height_arg = ['--export-height', str(height)] if height is not None else []
				dpi_arg = ['--export-dpi', str(dpi)] if dpi is not None else []
				if subprocess.run(['inkscape', full_path, '-o', target_path, *width_arg, *height_arg, *dpi_arg]).returncode != 0:
					raise Exception(f"Could not convert '{full_path}' to '{format}'")
			
			# Convert everything else using ImageMagick.
			else:
				resize_arg = ['-resize', str(geometry)] if geometry is not None else []
				density_arg = ['-density', str(dpi)] if dpi is not None else []
				quality_arg = ['-quality', str(quality)] if quality is not None else []
				if subprocess.run(['convert', *density_arg, full_path, *resize_arg, *quality_arg, target_path]).returncode != 0:
					raise Exception(f"Could not convert '{full_path}' to '{format}'")

		return target_name if relative else target_path


	def get_image_size(self, input_filename: str, additional_dirs: List[str]=[]) -> (int, int):
		full_path = self.find_image(input_filename, additional_dirs)
		if full_path is None:
			raise FileNotFoundError(f'Image {input_filename} not found.')
		# Getting image size using ImageMagick is slow. VERY
		return Image.open(full_path).size


	def find_image(self, input_filename: str, additional_dirs: List[str]=[]) -> str:
		for dir in [*self.lookup_dirs, *additional_dirs]:
			if os.path.isfile(dir + "/" + input_filename):
				return dir + "/" + input_filename