mirror of
https://github.com/squidfunk/mkdocs-material.git
synced 2026-06-03 22:12:07 -04:00
389 lines
15 KiB
Python
389 lines
15 KiB
Python
# Copyright (c) 2016-2025 Martin Donath <martin.donath@squidfunk.com>
|
||
|
||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||
# of this software and associated documentation files (the "Software"), to
|
||
# deal in the Software without restriction, including without limitation the
|
||
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
|
||
# sell copies of the Software, and to permit persons to whom the Software is
|
||
# furnished to do so, subject to the following conditions:
|
||
|
||
# The above copyright notice and this permission notice shall be included in
|
||
# all copies or substantial portions of the Software.
|
||
|
||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||
# FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
|
||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
|
||
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
|
||
# IN THE SOFTWARE.
|
||
|
||
from __future__ import annotations
|
||
|
||
import functools
|
||
import json
|
||
import logging
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
|
||
from fnmatch import fnmatch
|
||
from colorama import Fore, Style
|
||
from concurrent.futures import Future
|
||
from concurrent.futures.thread import ThreadPoolExecutor
|
||
from hashlib import sha1
|
||
from mkdocs import utils
|
||
from mkdocs.config.defaults import MkDocsConfig
|
||
from mkdocs.exceptions import PluginError
|
||
from mkdocs.plugins import BasePlugin
|
||
from mkdocs.structure.files import File
|
||
from shutil import which
|
||
try:
|
||
from PIL import Image
|
||
except ImportError:
|
||
pass
|
||
|
||
from .config import OptimizeConfig
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Classes
|
||
# -----------------------------------------------------------------------------
|
||
|
||
# Optimize plugin
|
||
class OptimizePlugin(BasePlugin[OptimizeConfig]):
|
||
supports_multiple_instances = True
|
||
|
||
# Manifest
|
||
manifest: dict[str, str] = {}
|
||
|
||
# Initialize plugin
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
|
||
# Initialize incremental builds
|
||
self.is_serve = False
|
||
|
||
# Determine whether we're serving the site
|
||
def on_startup(self, *, command, dirty):
|
||
self.is_serve = command == "serve"
|
||
|
||
# Initialize thread pool
|
||
self.pool = ThreadPoolExecutor(self.config.concurrency)
|
||
self.pool_jobs: dict[str, Future] = {}
|
||
|
||
# Resolve and load manifest
|
||
def on_config(self, config):
|
||
if not self.config.enabled:
|
||
return
|
||
|
||
# Resolve cache directory (once) - this is necessary, so the cache is
|
||
# always relative to the configuration file, and thus project, and not
|
||
# relative to the current working directory, or it would not work with
|
||
# the projects plugin.
|
||
path = os.path.abspath(self.config.cache_dir)
|
||
if path != self.config.cache_dir:
|
||
self.config.cache_dir = os.path.join(
|
||
os.path.dirname(config.config_file_path),
|
||
os.path.normpath(self.config.cache_dir)
|
||
)
|
||
|
||
# Ensure cache directory exists
|
||
os.makedirs(self.config.cache_dir, exist_ok = True)
|
||
|
||
# Initialize manifest
|
||
self.manifest_file = os.path.join(
|
||
self.config.cache_dir, "manifest.json"
|
||
)
|
||
|
||
# Load manifest if it exists and the cache should be used
|
||
if os.path.isfile(self.manifest_file) and self.config.cache:
|
||
try:
|
||
with open(self.manifest_file) as f:
|
||
self.manifest = json.load(f)
|
||
except:
|
||
pass
|
||
|
||
# Initialize optimization pipeline
|
||
def on_env(self, env, *, config, files):
|
||
if not self.config.enabled:
|
||
return
|
||
|
||
# Skip if media files should not be optimized
|
||
if not self.config.optimize:
|
||
return
|
||
|
||
# Filter all optimizable media files and steal reponsibility from MkDocs
|
||
# by removing them from the files collection. Then, start a concurrent
|
||
# job that checks if an image was already optimized and can be returned
|
||
# from the cache, or optimize it accordingly.
|
||
for file in files.media_files():
|
||
if self._is_excluded(file):
|
||
continue
|
||
|
||
# Spawn concurrent job to optimize the given image and add future
|
||
# to job dictionary, as it returns the file we need to copy later
|
||
path = os.path.join(self.config.cache_dir, file.src_path)
|
||
self.pool_jobs[file.abs_src_path] = self.pool.submit(
|
||
self._optimize_image, file, path, config
|
||
)
|
||
|
||
# Steal responsibility from MkDocs
|
||
files.remove(file)
|
||
|
||
# Finish optimization pipeline
|
||
def on_post_build(self, *, config):
|
||
if not self.config.enabled:
|
||
return
|
||
|
||
# Skip if media files should not be optimized
|
||
if not self.config.optimize:
|
||
return
|
||
|
||
# Reconcile concurrent jobs - we need to wait for all jobs to finish
|
||
# before we can copy the optimized files to the output directory. If an
|
||
# exception occurred in one of the jobs, we raise it here, so the build
|
||
# fails and the author can fix the issue.
|
||
for path, future in self.pool_jobs.items():
|
||
if future.exception():
|
||
raise future.exception()
|
||
else:
|
||
file: File = future.result()
|
||
file.copy_file()
|
||
|
||
# Save manifest if cache should be used
|
||
if self.config.cache:
|
||
with open(self.manifest_file, "w") as f:
|
||
f.write(json.dumps(self.manifest, indent = 2, sort_keys = True))
|
||
|
||
# Compute and print gains through optimization
|
||
if self.config.print_gain_summary:
|
||
print(Style.NORMAL)
|
||
print(f" Optimizations:")
|
||
|
||
# Print summary for file extension
|
||
for seek in [".png", ".jpg"]:
|
||
size = size_opt = 0
|
||
for path, future in self.pool_jobs.items():
|
||
file: File = future.result()
|
||
|
||
# Skip files that are not of the given type
|
||
_, extension = os.path.splitext(path)
|
||
extension = ".jpg" if extension == ".jpeg" else extension
|
||
if extension != seek:
|
||
continue
|
||
|
||
# Compute size before and after optimization
|
||
size += os.path.getsize(path)
|
||
size_opt += os.path.getsize(file.abs_dest_path)
|
||
|
||
# Compute absolute and relative gain
|
||
if size and size_opt:
|
||
gain_abs = size - size_opt
|
||
gain_rel = (1 - size_opt / size) * 100
|
||
|
||
# Print summary for files
|
||
print(
|
||
f" *{seek} {Fore.GREEN}{_size(size_opt)}"
|
||
f"{Fore.WHITE}{Style.DIM} ↓ "
|
||
f"{_size(gain_abs)} [{gain_rel:3.1f}%]"
|
||
f"{Style.RESET_ALL}"
|
||
)
|
||
|
||
# Reset all styles
|
||
print(Style.RESET_ALL)
|
||
|
||
# Save manifest on shutdown
|
||
def on_shutdown(self):
|
||
if not self.config.enabled:
|
||
return
|
||
|
||
# Shutdown thread pool - if we're on Python 3.9 and above, cancel all
|
||
# pending futures that have not yet been scheduled
|
||
if sys.version_info >= (3, 9):
|
||
self.pool.shutdown(cancel_futures = True)
|
||
else:
|
||
self.pool.shutdown()
|
||
|
||
# Save manifest if cache should be used
|
||
if self.manifest and self.config.cache:
|
||
with open(self.manifest_file, "w") as f:
|
||
f.write(json.dumps(self.manifest, indent = 2, sort_keys = True))
|
||
|
||
# -------------------------------------------------------------------------
|
||
|
||
# Check if a file can be optimized
|
||
def _is_optimizable(self, file: File):
|
||
|
||
# Check if PNG images should be optimized
|
||
if file.url.endswith((".png")):
|
||
return self.config.optimize_png
|
||
|
||
# Check if JPG images should be optimized
|
||
if file.url.endswith((".jpg", ".jpeg")):
|
||
return self.config.optimize_jpg
|
||
|
||
# File can not be optimized by the plugin
|
||
return False
|
||
|
||
# Check if the given file is excluded
|
||
def _is_excluded(self, file: File):
|
||
if not self._is_optimizable(file):
|
||
return True
|
||
|
||
# Check if file matches one of the inclusion patterns
|
||
path = file.src_path
|
||
if self.config.optimize_include:
|
||
for pattern in self.config.optimize_include:
|
||
if fnmatch(file.src_uri, pattern):
|
||
return False
|
||
|
||
# File is not included
|
||
log.debug(f"Excluding file '{path}' due to inclusion patterns")
|
||
return True
|
||
|
||
# Check if file matches one of the exclusion patterns
|
||
for pattern in self.config.optimize_exclude:
|
||
if fnmatch(file.src_uri, pattern):
|
||
log.debug(f"Excluding file '{path}' due to exclusion patterns")
|
||
return True
|
||
|
||
# File is not excluded
|
||
return False
|
||
|
||
# Optimize image and write to cache
|
||
def _optimize_image(self, file: File, path: str, config: MkDocsConfig):
|
||
with open(file.abs_src_path, "rb") as f:
|
||
data = f.read()
|
||
hash = sha1(data).hexdigest()
|
||
|
||
# Check if file hash changed, so we need to optimize again
|
||
prev = self.manifest.get(file.url, "")
|
||
if hash != prev or not os.path.isfile(path):
|
||
os.makedirs(os.path.dirname(path), exist_ok = True)
|
||
|
||
# Optimize PNG image using pngquant
|
||
if file.url.endswith((".png")):
|
||
self._optimize_image_png(file, path, config)
|
||
|
||
# Optimize JPG image using pillow
|
||
if file.url.endswith((".jpg", ".jpeg")):
|
||
self._optimize_image_jpg(file, path, config)
|
||
|
||
# Compute size before and after optimization
|
||
size = len(data)
|
||
size_opt = os.path.getsize(path)
|
||
|
||
# Compute absolute and relative gain
|
||
gain_abs = size - size_opt
|
||
gain_rel = (1 - size_opt / size) * 100
|
||
|
||
# Print how much we gained, if we did and desired
|
||
gain = ""
|
||
if gain_abs and self.config.print_gain:
|
||
gain += " ↓ "
|
||
gain += " ".join([_size(gain_abs), f"[{gain_rel:3.1f}%]"])
|
||
|
||
# Print summary for file
|
||
log.info(
|
||
f"Optimized media file: {file.src_uri} "
|
||
f"{Fore.GREEN}{_size(size_opt)}"
|
||
f"{Fore.WHITE}{Style.DIM}{gain}"
|
||
f"{Style.RESET_ALL}"
|
||
)
|
||
|
||
# Update manifest by associating file with hash
|
||
self.manifest[file.url] = hash
|
||
|
||
# Compute project root
|
||
root = os.path.dirname(config.config_file_path)
|
||
|
||
# Compute source file system path
|
||
file.abs_src_path = path
|
||
file.src_path = os.path.relpath(path, root)
|
||
|
||
# Return file to be copied from cache
|
||
return file
|
||
|
||
# Optimize PNG image - we first tried to use libimagequant, but encountered
|
||
# the occassional segmentation fault, which means it's probably not a good
|
||
# choice. Instead, we just rely on pngquant which seems much more stable.
|
||
def _optimize_image_png(self, file: File, path: str, config: MkDocsConfig):
|
||
|
||
# Check if the required dependencies for optimizing are available, which
|
||
# is, at the absolute minimum, the 'pngquant' binary, and raise an error
|
||
# to the caller, so he can decide what to do with the error. The caller
|
||
# can treat this as a warning or an error to abort the build.
|
||
if not which("pngquant"):
|
||
docs = os.path.relpath(config.docs_dir)
|
||
path = os.path.relpath(file.abs_src_path, docs)
|
||
raise PluginError(
|
||
f"Couldn't optimize image '{path}' in '{docs}': 'pngquant' "
|
||
f"not found. Make sure 'pngquant' is installed and in your path"
|
||
)
|
||
|
||
# Build command line arguments
|
||
args = ["pngquant",
|
||
"--force", "--skip-if-larger",
|
||
"--output", path,
|
||
"--speed", f"{self.config.optimize_png_speed}"
|
||
]
|
||
|
||
# Add flag to remove optional metadata
|
||
if self.config.optimize_png_strip:
|
||
args.append("--strip")
|
||
|
||
# Set input file and run, then check if pngquant actually wrote a file,
|
||
# as we instruct it not to if the size of the optimized file is larger.
|
||
# This can happen if files are already compressed and optimized by
|
||
# the author. In that case, just copy the original file.
|
||
subprocess.run([*args, file.abs_src_path])
|
||
if not os.path.isfile(path):
|
||
utils.copy_file(file.abs_src_path, path)
|
||
|
||
# Optimize JPG image
|
||
def _optimize_image_jpg(self, file: File, path: str, config: MkDocsConfig):
|
||
|
||
# Check if the required dependencies for optimizing are available, which
|
||
# is, at the absolute minimum, the 'pillow' package, and raise an error
|
||
# to the caller, so he can decide what to do with the error. The caller
|
||
# can treat this as a warning or an error to abort the build.
|
||
if not _supports("Image"):
|
||
docs = os.path.relpath(config.docs_dir)
|
||
path = os.path.relpath(file.abs_src_path, docs)
|
||
raise PluginError(
|
||
f"Couldn't optimize image '{path}' in '{docs}': install "
|
||
f"required dependencies – pip install 'mkdocs-material[imaging]'"
|
||
)
|
||
|
||
# Open and save optimized image
|
||
image = Image.open(file.abs_src_path)
|
||
image.save(path, "jpeg",
|
||
quality = self.config.optimize_jpg_quality,
|
||
progressive = self.config.optimize_jpg_progressive
|
||
)
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Helper functions
|
||
# -----------------------------------------------------------------------------
|
||
|
||
# Check for presence of optional imports
|
||
@functools.lru_cache(maxsize = None)
|
||
def _supports(name: str):
|
||
return name in globals()
|
||
|
||
# -----------------------------------------------------------------------------
|
||
|
||
# Print human-readable size
|
||
def _size(value):
|
||
for unit in ["B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB"]:
|
||
if abs(value) < 1000.0:
|
||
return f"{value:3.1f} {unit}"
|
||
value /= 1000.0
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# Data
|
||
# -----------------------------------------------------------------------------
|
||
|
||
# Set up logging
|
||
log = logging.getLogger("mkdocs.material.optimize")
|