Explicit caching for Prawn PDF images

prawn-grouping.rb

# frozen_string_literal: true

module Prawn
  # `cached_image` needs to know if it's in a cloned document or the original so that it can store
  # the image object data properly.
  class ClonedDocument < Document
    attr_reader :root_document

    def initialize(root_document, **options, &block)
      @root_document = root_document

      super(**options, &block)
    end

    protected

    # Make sure the image cache used is the parent document's cache, not a temporary
    # one that'll be lost after the group ends.
    def cached_images
      @root_document.cached_images
    end
  end

  # This is mostly just a rewrite of `prawn-grouping` to fit my needs. The only necessary change is
  # that it creates a ClonedDocument instead of a Document.
  module Grouping
    def group
      # create a temporary document with current context and offset
      cloned_pdf = create_box_clone
      update_cloned_pdf cloned_pdf

      yield cloned_pdf

      start_new_page if cloned_pdf.page_count > 1

      yield self
    end

    private

    def create_box_clone
      Prawn::ClonedDocument.new(
        self,
        page_size: state.page.size,
        page_layout: state.page.layout
      )
    end

    def update_cloned_pdf(cloned)
      cloned.margin_box = @bounding_box.dup
      cloned.text_formatter = @text_formatter.dup
      # cloned.page.margins = page.margins.dup
      cloned.font_families.merge! font_families
      cloned.font font.family
      cloned.font_size font_size
      cloned.default_leading = default_leading
      cloned.y = y
    end
  end
end

Prawn::Document.extensions << Prawn::Grouping

prawn-images.rb

# frozen_string_literal: true

module Prawn
  module Images
    def cached_image(cache_key, **options)
      Prawn.verify_options(%i[at position vposition height width scale fit], options)

      unless cached_images.key?(cache_key)
        raise ArgumentError, "Cached image #{cache_key} not found" unless block_given?

        yield(->(file) { cached_images[cache_key] = cacheable_image_data(file) })
      end

      pdf_obj, info = cached_images[cache_key]

      embed_image(pdf_obj, info, options)

      info
    end

    # This is just `image` with all of the caching logic bypassed.
    def uncached_image(file, **options)
      Prawn.verify_options(%i[at position vposition height width scale fit], options)

      pdf_obj, info = cacheable_image_data(file)

      embed_image(pdf_obj, info, options)

      info
    end

    protected

    def cached_images
      @cached_images ||= {}
    end

    # This doesn't *perform* caching, it's just data that *can be* cached. It's basically just the
    # contents of `build_image_object`'s `else` branch.
    def cacheable_image_data(file)
      image_content = verify_and_read_image(file)

      # Build the image object
      info = Prawn.image_handler.find(image_content).new(image_content)

      # Bump PDF version if the image requires it
      renderer.min_version(info.min_pdf_version) if info.respond_to?(:min_pdf_version)

      # Add the image to the PDF and register it in case we see it again. This must be given
      # the original PDF, not a ClonedDocument, or the actual image data will be lost before
      # it's actually used.
      image_obj = info.build_pdf_object(respond_to?(:root_document) ? root_document : self)

      [image_obj, info]
    end
  end
end

usage.rb

# frozen_string_literal: true

# Inside the prawn template:

# An IO/File/string can be used, just like with the original `image` method. This method will only
# read the file from disk once.
pdf.cached_image('vendor_logo', scale: 0.25) do |block|
  block.call 'path/to/vendor_logo.png'
end

# Active Storage attachments need to be opened first and passed the yielded block
pdf.cached_image("user_avatar_#{user.id}", height: 50) do |block|
  user.avatar.open(&block)
end

# This will never be cached, and will be read from disk every time.
pdf.uncached_image('some/args/as/image.png', width: 180)