pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/utils/description_helper.py

@@ -1,5 +1,4 @@
 import dataclasses
-from typing import Optional, Union
 
 import pandas as pd
 from pandas.io.formats.style import Styler
@@ -7,11 +6,11 @@ from pandas.io.formats.style import Styler
 
 @dataclasses.dataclass
 class _Descriptor:
-    body: Union[str, pd.DataFrame]
+    body: str | pd.DataFrame
     # The remaining fields only affect the behavior if `body` is a pd.DataFrame.
     show_index: bool
     show_header: bool
-    styler: Optional[Styler] = None
+    styler: Styler | None = None
 
 
 class DescriptionHelper:
@@ -25,6 +24,7 @@ class DescriptionHelper:
     DescriptionHelper can convert a list of descriptors into either HTML or plaintext and do something reasonable
     in each case.
     """
+
     __descriptors: list[_Descriptor]
 
     def __init__(self) -> None:
@@ -32,10 +32,10 @@ class DescriptionHelper:
 
     def append(
         self,
-        descriptor: Union[str, pd.DataFrame],
+        descriptor: str | pd.DataFrame,
         show_index: bool = False,
         show_header: bool = True,
-        styler: Optional[Styler] = None,
+        styler: Styler | None = None,
     ) -> None:
         self.__descriptors.append(_Descriptor(descriptor, show_index, show_header, styler))
 
@@ -69,18 +69,17 @@ class DescriptionHelper:
             return (
                 # Render the string as a single-cell DataFrame. This will ensure a consistent style of output in
                 # cases where strings appear alongside DataFrames in the same DescriptionHelper.
-                pd.DataFrame([descriptor.body]).style
-                .set_properties(None, **{'white-space': 'pre-wrap', 'text-align': 'left', 'font-weight': 'bold'})
-                .hide(axis='index').hide(axis='columns')
+                pd.DataFrame([descriptor.body])
+                .style.set_properties(None, **{'white-space': 'pre-wrap', 'text-align': 'left', 'font-weight': 'bold'})
+                .hide(axis='index')
+                .hide(axis='columns')
             )
         else:
             styler = descriptor.styler
             if styler is None:
                 styler = descriptor.body.style
-            styler = (
-                styler
-                .set_properties(None, **{'white-space': 'pre-wrap', 'text-align': 'left'})
-                .set_table_styles([dict(selector='th', props=[('text-align', 'left')])])
+            styler = styler.set_properties(None, **{'white-space': 'pre-wrap', 'text-align': 'left'}).set_table_styles(
+                [{'selector': 'th', 'props': [('text-align', 'left')]}]
             )
             if not descriptor.show_header:
                 styler = styler.hide(axis='columns')

pixeltable/utils/documents.py

@@ -1,86 +1,85 @@
 import dataclasses
-from typing import Optional
+import os
 
 import bs4
-import fitz  # type: ignore[import-untyped]
 import puremagic
+from pypdfium2 import PdfDocument  # type: ignore[import-untyped]
 
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 from pixeltable.env import Env
 
 
 @dataclasses.dataclass
 class DocumentHandle:
     format: ts.DocumentType.DocumentFormat
-    bs_doc: Optional[bs4.BeautifulSoup] = None
-    md_ast: Optional[dict] = None
-    pdf_doc: Optional[fitz.Document] = None
+    bs_doc: bs4.BeautifulSoup | None = None
+    md_ast: dict | None = None
+    pdf_doc: PdfDocument | None = None
+    txt_doc: str | None = None
 
 
-def get_document_handle(path: str) -> Optional[DocumentHandle]:
-    doc_format = puremagic.from_file(path)
+def get_document_handle(path: str) -> DocumentHandle:
+    _, extension = os.path.splitext(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.pdf':
-        pdf_doc = get_pdf_handle(path)
-        if pdf_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.PDF, pdf_doc=pdf_doc)
+    # if no extension, use puremagic to determine the type
+    extension = puremagic.from_file(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.html':
-        bs_doc = get_html_handle(path)
-        if bs_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.HTML, bs_doc=bs_doc)
+    raise excs.Error(f'Unrecognized document format: {path}')
 
-    if doc_format == '.md':
-        md_ast = get_markdown_handle(path)
-        if md_ast is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.MD, md_ast=md_ast)
 
-    if doc_format == '.xml':
-        bs_doc = get_xml_handle(path)
-        if bs_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.XML, bs_doc=bs_doc)
+def get_handle_by_extension(path: str, extension: str) -> DocumentHandle | None:
+    doc_format = ts.DocumentType.DocumentFormat.from_extension(extension)
 
-    return None
-
-
-def get_pdf_handle(path: str) -> Optional[fitz.Document]:
-    try:
-        doc = fitz.open(path)
-        # check pdf (bc it will work for images)
-        if not doc.is_pdf:
-            return None
-        # try to read one page
-        next(page for page in doc)
-        return doc
-    except Exception:
-        return None
-
-
-def get_html_handle(path: str) -> Optional[bs4.BeautifulSoup]:
     try:
-        with open(path, 'r', encoding='utf8') as fp:
-            doc = bs4.BeautifulSoup(fp, 'lxml')
-        return doc if doc.find() is not None else None
-    except Exception:
-        return None
+        if doc_format == ts.DocumentType.DocumentFormat.HTML:
+            return DocumentHandle(doc_format, bs_doc=get_html_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.MD:
+            return DocumentHandle(doc_format, md_ast=get_markdown_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.PDF:
+            return DocumentHandle(doc_format, pdf_doc=PdfDocument(path))
+        if doc_format == ts.DocumentType.DocumentFormat.XML:
+            return DocumentHandle(doc_format, bs_doc=get_xml_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.TXT:
+            return DocumentHandle(doc_format, txt_doc=get_txt(path))
+    except Exception as exc:
+        raise excs.Error(f'An error occurred processing a {doc_format} document: {path}') from exc
 
+    return None
 
-def get_xml_handle(path: str) -> Optional[bs4.BeautifulSoup]:
-    try:
-        with open(path, 'r', encoding='utf8') as fp:
-            doc = bs4.BeautifulSoup(fp, 'xml')
-        return doc if doc.find() is not None else None
-    except Exception:
-        return None
+
+def get_html_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'lxml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid HTML document: {path}')
+    return doc
 
 
-def get_markdown_handle(path: str) -> Optional[dict]:
+def get_markdown_handle(path: str) -> dict:
     Env.get().require_package('mistune', [3, 0])
     import mistune
-    try:
-        with open(path, encoding='utf8') as file:
-            text = file.read()
-        md_ast = mistune.create_markdown(renderer=None)
-        return md_ast(text)
-    except Exception:
-        return None
+
+    with open(path, encoding='utf8') as file:
+        text = file.read()
+    md_ast = mistune.create_markdown(renderer=None)
+    return md_ast(text)
+
+
+def get_xml_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'xml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid XML document: {path}')
+    return doc
+
+
+def get_txt(path: str) -> str:
+    with open(path, 'r', encoding='utf-8') as fp:
+        doc = fp.read()
+    return doc

pixeltable/utils/exception_handler.py

@@ -0,0 +1,36 @@
+import logging
+from typing import Any, Callable, TypeVar
+
+R = TypeVar('R')
+
+logger = logging.getLogger('pixeltable')
+
+
+def run_cleanup(cleanup_func: Callable[..., R], *args: Any, raise_error: bool = True, **kwargs: Any) -> R | None:
+    """
+    Runs a cleanup function. If interrupted, retry cleanup.
+    The `run_cleanup()` function ensures that the `cleanup_func()` function executes at least once.
+    If the `cleanup_func()` is interrupted during execution, it will be retried.
+
+    Args:
+        cleanup_func: an idempotent function
+        raise_error: raise an exception if an error occurs during cleanup.
+    """
+    try:
+        logger.debug(f'Running cleanup function: {cleanup_func.__name__!r}')
+        return cleanup_func(*args, **kwargs)
+    except KeyboardInterrupt as interrupt:
+        # Save original exception and re-attempt cleanup
+        original_exception = interrupt
+        logger.debug(f'Cleanup {cleanup_func.__name__!r} interrupted, retrying')
+        try:
+            return cleanup_func(*args, **kwargs)
+        except Exception as e:
+            # Suppress this exception
+            logger.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e.__class__}: {e}')
+        raise KeyboardInterrupt from original_exception
+    except Exception as e:
+        logger.error(f'Cleanup {cleanup_func.__name__!r} failed with exception {e.__class__}: {e}')
+        if raise_error:
+            raise e
+    return None

pixeltable/utils/filecache.py

@@ -5,21 +5,22 @@ import hashlib
 import logging
 import os
 import warnings
-from collections import OrderedDict, defaultdict, namedtuple
+from collections import OrderedDict, defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Optional
+from typing import NamedTuple
 from uuid import UUID
 
 import pixeltable.exceptions as excs
+from pixeltable.config import Config
 from pixeltable.env import Env
 
 _logger = logging.getLogger('pixeltable')
 
+
 @dataclass
 class CacheEntry:
-
     key: str
     tbl_id: UUID
     col_id: int
@@ -56,7 +57,8 @@ class FileCache:
     TODO:
     - implement MRU eviction for queries that exceed the capacity
     """
-    __instance: Optional[FileCache] = None
+
+    __instance: FileCache | None = None
 
     cache: OrderedDict[str, CacheEntry]
     total_size: int
@@ -77,11 +79,18 @@ class FileCache:
     evicted_working_set_keys: set[str]
     new_redownload_witnessed: bool  # whether a new re-download has occurred since the last time a warning was issued
 
-    FileCacheColumnStats = namedtuple('FileCacheColumnStats', ('tbl_id', 'col_id', 'num_files', 'total_size'))
-    FileCacheStats = namedtuple(
-        'FileCacheStats',
-        ('total_size', 'num_requests', 'num_hits', 'num_evictions', 'column_stats')
-    )
+    class FileCacheColumnStats(NamedTuple):
+        tbl_id: UUID
+        col_id: int
+        num_files: int
+        total_size: int
+
+    class FileCacheStats(NamedTuple):
+        total_size: int
+        num_requests: int
+        num_hits: int
+        num_evictions: int
+        column_stats: list[FileCache.FileCacheColumnStats]
 
     @classmethod
     def get(cls) -> FileCache:
@@ -93,7 +102,7 @@ class FileCache:
     def init(cls) -> None:
         cls.__instance = cls()
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.cache = OrderedDict()
         self.total_size = 0
         self.capacity_bytes = int(Env.get()._file_cache_size_g * (1 << 30))
@@ -117,17 +126,18 @@ class FileCache:
             return 0
         return int(self.total_size / len(self.cache))
 
-    def num_files(self, tbl_id: Optional[UUID] = None) -> int:
+    def num_files(self, tbl_id: UUID | None = None) -> int:
         if tbl_id is None:
             return len(self.cache)
         return sum(e.tbl_id == tbl_id for e in self.cache.values())
 
-    def clear(self, tbl_id: Optional[UUID] = None) -> None:
+    def clear(self, tbl_id: UUID | None = None) -> None:
         """
         For testing purposes: allow resetting capacity and stats.
         """
         if tbl_id is None:
-            # We need to store the entries to remove in a list, because we can't remove items from a dict while iterating
+            # We need to store the entries to remove in a list, because we can't remove items from a dict
+            # while iterating
             entries_to_remove = list(self.cache.values())
             _logger.debug(f'clearing {self.num_files()} entries from file cache')
             self.num_requests, self.num_hits, self.num_evictions = 0, 0, 0
@@ -153,8 +163,9 @@ class FileCache:
                 f'of the evicted file(s) is {round(extra_capacity_needed / (1 << 30), 1)} GiB.\n'
                 f'Consider increasing the cache size to at least {round(suggested_cache_size / (1 << 30), 1)} GiB '
                 f'(it is currently {round(self.capacity_bytes / (1 << 30), 1)} GiB).\n'
-                f'You can do this by setting the value of `file_cache_size_g` in: {str(Env.get()._config_file)}',
-                excs.PixeltableWarning
+                f'You can do this by setting the value of `file_cache_size_g` in: {Config.get().config_file}',
+                excs.PixeltableWarning,
+                stacklevel=2,
             )
             self.new_redownload_witnessed = False
 
@@ -163,7 +174,7 @@ class FileCache:
         h.update(url.encode())
         return h.hexdigest()
 
-    def lookup(self, url: str) -> Optional[Path]:
+    def lookup(self, url: str) -> Path | None:
         self.num_requests += 1
         key = self._url_hash(url)
         entry = self.cache.get(key, None)
@@ -195,13 +206,15 @@ class FileCache:
             self.evicted_working_set_keys.add(key)
             self.new_redownload_witnessed = True
         self.keys_retrieved.add(key)
-        entry = CacheEntry(key, tbl_id, col_id, file_info.st_size, datetime.fromtimestamp(file_info.st_mtime), path.suffix)
+        entry = CacheEntry(
+            key, tbl_id, col_id, file_info.st_size, datetime.fromtimestamp(file_info.st_mtime), path.suffix
+        )
         self.cache[key] = entry
         self.total_size += entry.size
         new_path = entry.path
         os.rename(str(path), str(new_path))
         new_path.touch(exist_ok=True)
-        _logger.debug(f'added entry for cell {url} to file cache')
+        _logger.debug(f'FileCache: cached url {url} with file name {new_path}')
         return new_path
 
     def ensure_capacity(self, size: int) -> None:
@@ -217,7 +230,9 @@ class FileCache:
                 # Make a record of the eviction, so that we can generate a warning later if the key is retrieved again.
                 self.keys_evicted_after_retrieval.add(lru_entry.key)
             os.remove(str(lru_entry.path))
-            _logger.debug(f'evicted entry for cell {lru_entry.key} from file cache (of size {lru_entry.size // (1 << 20)} MiB)')
+            _logger.debug(
+                f'evicted entry for cell {lru_entry.key} from file cache (of size {lru_entry.size // (1 << 20)} MiB)'
+            )
 
     def set_capacity(self, capacity_bytes: int) -> None:
         self.capacity_bytes = capacity_bytes
@@ -228,15 +243,16 @@ class FileCache:
         # (tbl_id, col_id) -> (num_files, total_size)
         d: dict[tuple[UUID, int], list[int]] = defaultdict(lambda: [0, 0])
         for entry in self.cache.values():
-            t = d[(entry.tbl_id, entry.col_id)]
+            t = d[entry.tbl_id, entry.col_id]
             t[0] += 1
             t[1] += entry.size
         col_stats = [
-            self.FileCacheColumnStats(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()
+            self.FileCacheColumnStats(tbl_id, col_id, num_files, size)
+            for (tbl_id, col_id), (num_files, size) in d.items()
         ]
         col_stats.sort(key=lambda e: e[3], reverse=True)
         return self.FileCacheStats(self.total_size, self.num_requests, self.num_hits, self.num_evictions, col_stats)
 
     def debug_print(self) -> None:
         for entry in self.cache.values():
-            print(f'CacheEntry: tbl_id={entry.tbl_id}, col_id={entry.col_id}, size={entry.size}')
+            _logger.debug(f'CacheEntry: tbl_id={entry.tbl_id}, col_id={entry.col_id}, size={entry.size}')

pixeltable/utils/formatter.py

@@ -4,12 +4,13 @@ import io
 import json
 import logging
 import mimetypes
-from typing import Any, Callable, Optional
+import uuid
+from typing import Any, Callable
 
-import av  # type: ignore[import-untyped]
+import av
 import numpy as np
-import PIL
-import PIL.Image as Image
+from PIL import Image
+from pypdfium2 import PdfDocument  # type: ignore[import-untyped]
 
 import pixeltable.type_system as ts
 from pixeltable.utils.http_server import get_file_uri
@@ -20,11 +21,11 @@ _logger = logging.getLogger('pixeltable')
 class Formatter:
     """
     A factory for constructing HTML formatters for Pixeltable data. The formatters are used to customize
-    the rendering of `DataFrameResultSet`s in notebooks.
+    the rendering of `ResultSet`s in notebooks.
 
     Args:
-        num_rows: Number of rows in the DataFrame being rendered.
-        num_cols: Number of columns in the DataFrame being rendered.
+        num_rows: Number of rows in the `ResultSet` being rendered.
+        num_cols: Number of columns in the `ResultSet` being rendered.
         http_address: Root address of the Pixeltable HTTP server (used to construct URLs for media references).
     """
 
@@ -40,9 +41,13 @@ class Formatter:
         self.__num_cols = num_cols
         self.__http_address = http_address
 
-    def get_pandas_formatter(self, col_type: ts.ColumnType) -> Optional[Callable]:
+    def get_pandas_formatter(self, col_type: ts.ColumnType) -> Callable | None:
         if col_type.is_string_type():
             return self.format_string
+        if col_type.is_uuid_type():
+            return self.format_uuid
+        if col_type.is_binary_type():
+            return self.format_binary
         if col_type.is_float_type():
             return self.format_float
         if col_type.is_json_type():
@@ -64,10 +69,24 @@ class Formatter:
         """
         Escapes special characters in `val`, and abbreviates `val` if its length exceeds `_STRING_MAX_LEN`.
         """
-        return cls.__escape(cls.__abbreviate(val, cls.__STRING_MAX_LEN))
+        return cls.__escape(cls.abbreviate(val))
 
     @classmethod
-    def __abbreviate(cls, val: str, max_len: int) -> str:
+    def format_uuid(cls, val: uuid.UUID | None) -> str:
+        """
+        Formats a UUID by converting it to a string and applying string formatting.
+        """
+        return '' if val is None else cls.format_string(str(val))
+
+    @classmethod
+    def format_binary(cls, val: bytes) -> str:
+        """
+        Formats binary data by converting it to an encoded string and applying string formatting.
+        """
+        return cls.format_string(str(val))
+
+    @classmethod
+    def abbreviate(cls, val: str, max_len: int = __STRING_MAX_LEN) -> str:
         if len(val) > max_len:
             edgeitems = (max_len - len(cls.__STRING_SEP)) // 2
             return f'{val[:edgeitems]}{cls.__STRING_SEP}{val[-edgeitems:]}'
@@ -95,41 +114,45 @@ class Formatter:
         )
 
     @classmethod
-    def format_json(cls, val: Any) -> str:
+    def format_json(cls, val: Any, escape_strings: bool = True) -> str:
         if isinstance(val, str):
             # JSON-like formatting will be applied to strings that appear nested within a list or dict
             # (quote the string; escape any quotes inside the string; shorter abbreviations).
             # However, if the string appears in top-level position (i.e., the entire JSON value is a
             # string), then we format it like an ordinary string.
-            return cls.format_string(val)
+            return cls.format_string(val) if escape_strings else cls.abbreviate(val)
         # In all other cases, dump the JSON struct recursively.
-        return cls.__format_json_rec(val)
+        return cls.__format_json_rec(val, escape_strings)
 
     @classmethod
-    def __format_json_rec(cls, val: Any) -> str:
+    def __format_json_rec(cls, val: Any, escape_strings: bool) -> str:
         if isinstance(val, str):
-            return cls.__escape(json.dumps(cls.__abbreviate(val, cls.__NESTED_STRING_MAX_LEN)))
+            formatted = json.dumps(cls.abbreviate(val, cls.__NESTED_STRING_MAX_LEN))
+            return cls.__escape(formatted) if escape_strings else formatted
         if isinstance(val, float):
             return cls.format_float(val)
         if isinstance(val, np.ndarray):
             return cls.format_array(val)
         if isinstance(val, list):
             if len(val) < cls.__LIST_THRESHOLD:
-                components = [cls.__format_json_rec(x) for x in val]
+                components = [cls.__format_json_rec(x, escape_strings) for x in val]
             else:
-                components = [cls.__format_json_rec(x) for x in val[: cls.__LIST_EDGEITEMS]]
+                components = [cls.__format_json_rec(x, escape_strings) for x in val[: cls.__LIST_EDGEITEMS]]
                 components.append('...')
-                components.extend(cls.__format_json_rec(x) for x in val[-cls.__LIST_EDGEITEMS :])
+                components.extend(cls.__format_json_rec(x, escape_strings) for x in val[-cls.__LIST_EDGEITEMS :])
             return '[' + ', '.join(components) + ']'
         if isinstance(val, dict):
-            kv_pairs = (f'{cls.__format_json_rec(k)}: {cls.__format_json_rec(v)}' for k, v in val.items())
+            kv_pairs = (
+                f'{cls.__format_json_rec(k, escape_strings)}: {cls.__format_json_rec(v, escape_strings)}'
+                for k, v in val.items()
+            )
             return '{' + ', '.join(kv_pairs) + '}'
 
         # Everything else
         try:
             return json.dumps(val)
         except TypeError:  # Not JSON serializable
-            return str(val)
+            return cls.__escape(str(val))
 
     def format_img(self, img: Image.Image) -> str:
         """
@@ -153,22 +176,19 @@ class Formatter:
             """
 
     def format_video(self, file_path: str) -> str:
-        thumb_tag = ''
         # Attempt to extract the first frame of the video to use as a thumbnail,
         # so that the notebook can be exported as HTML and viewed in contexts where
         # the video itself is not accessible.
         # TODO(aaron-siegel): If the video is backed by a concrete external URL,
         # should we link to that instead?
-        with av.open(file_path) as container:
-            try:
-                thumb = next(container.decode(video=0)).to_image()
-                assert isinstance(thumb, Image.Image)
-                with io.BytesIO() as buffer:
-                    thumb.save(buffer, 'jpeg')
-                    thumb_base64 = base64.b64encode(buffer.getvalue()).decode()
-                    thumb_tag = f'poster="data:image/jpeg;base64,{thumb_base64}"'
-            except Exception:
-                pass
+        thumb = self.extract_first_video_frame(file_path)
+        if thumb is None:
+            thumb_tag = ''
+        else:
+            with io.BytesIO() as buffer:
+                thumb.save(buffer, 'jpeg')
+                thumb_base64 = base64.b64encode(buffer.getvalue()).decode()
+                thumb_tag = f'poster="data:image/jpeg;base64,{thumb_base64}"'
         if self.__num_rows > 1:
             width = 320
         elif self.__num_cols > 1:
@@ -183,6 +203,16 @@ class Formatter:
         </div>
         """
 
+    @classmethod
+    def extract_first_video_frame(cls, file_path: str) -> Image.Image | None:
+        with av.open(file_path) as container:
+            try:
+                img = next(container.decode(video=0)).to_image()
+                assert isinstance(img, Image.Image)
+                return img
+            except Exception:
+                return None
+
     def format_audio(self, file_path: str) -> str:
         return f"""
         <div class="pxt_audio">
@@ -192,29 +222,18 @@ class Formatter:
         </div>
         """
 
-    def format_document(self, file_path: str) -> str:
-        max_width = max_height = 320
+    def format_document(self, file_path: str, max_width: int = 320, max_height: int = 320) -> str:
         # by default, file path will be shown as a link
         inner_element = file_path
         inner_element = html.escape(inner_element)
-        # try generating a thumbnail for different types and use that if successful
-        if file_path.lower().endswith('.pdf'):
-            try:
-                import fitz  # type: ignore[import-untyped]
-
-                doc = fitz.open(file_path)
-                p = doc.get_page_pixmap(0)
-                while p.width > max_width or p.height > max_height:
-                    # shrink(1) will halve each dimension
-                    p.shrink(1)
-                data = p.tobytes(output='jpeg')
-                thumb_base64 = base64.b64encode(data).decode()
-                img_src = f'data:image/jpeg;base64,{thumb_base64}'
-                inner_element = f"""
-                    <img style="object-fit: contain; border: 1px solid black;" src="{img_src}" />
-                """
-            except:
-                logging.warning(f'Failed to produce PDF thumbnail {file_path}. Make sure you have PyMuPDF installed.')
+
+        thumb = self.make_document_thumbnail(file_path, max_width, max_height)
+        if thumb is not None:
+            with io.BytesIO() as buffer:
+                thumb.save(buffer, 'webp')
+                thumb_base64 = base64.b64encode(buffer.getvalue()).decode()
+                thumb_tag = f'data:image/webp;base64,{thumb_base64}'
+            inner_element = f'<img style="object-fit: contain; border: 1px solid black;" src="{thumb_tag}" />'
 
         return f"""
         <div class="pxt_document" style="width:{max_width}px;">
@@ -224,6 +243,24 @@ class Formatter:
         </div>
         """
 
+    @classmethod
+    def make_document_thumbnail(cls, file_path: str, max_width: int = 320, max_height: int = 320) -> Image.Image | None:
+        """
+        Returns a thumbnail image of a document.
+        """
+        if file_path.lower().endswith('.pdf'):
+            try:
+                doc = PdfDocument(file_path)
+                if len(doc) == 0:
+                    return None
+                img = doc[0].render().to_pil()
+                img.thumbnail((max_width, max_height), Image.LANCZOS)
+                return img
+            except Exception:
+                logging.warning(f'Failed to produce PDF thumbnail {file_path}. Make sure you have pypdfium2 installed.')
+
+        return None
+
     @classmethod
     def __create_source_tag(cls, http_address: str, file_path: str) -> str:
         src_url = get_file_uri(http_address, file_path)