from __future__ import annotations
import io
import pathlib
import zipfile
import wrapt
@wrapt.patch_function_wrapper('mimetypes', 'init')
def patch_mimetypes_init(wrapped, instance, args, kwargs):
"""Patch mimetypes init to use custom known files."""
knownfiles = pathlib.Path(__file__).parent.absolute() / 'mime.types'
return wrapped([str(knownfiles)])
import mimetypes
# Zip-container documents that libmagic reports as the generic
# ``application/zip``: OOXML (docx/xlsx/pptx) carry a ``[Content_Types].xml``
# alongside a format-specific top-level part directory; ODF stores its type
# verbatim in an uncompressed ``mimetype`` member.
_OOXML_MIME_BY_TOPDIR = {
'word': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
'xl': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'ppt': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
}
# What libmagic returns when it cannot see past the zip container.
_ZIP_CONTAINER_MIMES = frozenset({
'application/zip', 'application/octet-stream'})
# Office / OOXML / ODF content types that minimal base images (notably the
# AWS Lambda python base) carry in neither the builtin mimetypes map nor a
# system mime.types, so guess_extension(mime) returns None there. Register
# the canonical mime->extension maps so resolution is consistent on any
# image without relying on a system database.
_DOCUMENT_MIMETYPES = {
'application/msword': '.doc',
'application/vnd.ms-excel': '.xls',
'application/vnd.ms-powerpoint': '.ppt',
'application/vnd.ms-excel.sheet.binary.macroEnabled.12': '.xlsb',
'application/vnd.ms-excel.sheet.macroEnabled.12': '.xlsm',
'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
'.docx',
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
'.xlsx',
'application/vnd.openxmlformats-officedocument.presentationml.presentation':
'.pptx',
'application/vnd.oasis.opendocument.spreadsheet': '.ods',
'application/onenote': '.one',
}
def _register_document_mimetypes() -> None:
"""Register office/OOXML/ODF mime->extension maps with ``mimetypes``.
Idempotent. Runs at import so ``guess_extension`` resolves these types
on base images lacking a system mime.types (e.g. AWS Lambda). The mime
is lowercased because ``guess_extension`` lowercases its lookup key but
``add_type`` stores the type verbatim -- without this the mixed-case
``macroEnabled.12`` types would never resolve.
"""
for mime, ext in _DOCUMENT_MIMETYPES.items():
mimetypes.add_type(mime.lower(), ext)
_register_document_mimetypes()
def _refine_zip_mime(buffer: bytes) -> str | None:
"""Resolve a zip-container buffer to its precise OOXML/ODF mimetype.
libmagic reports OOXML (docx/xlsx/pptx) and many ODF documents as the
generic ``application/zip``. Inspect the archive to recover the true
type: ODF from its verbatim ``mimetype`` member, OOXML from the
top-level part directory next to ``[Content_Types].xml``. The whole
file is required (the zip directory lives at the end of the stream).
:param bytes buffer: Full file bytes.
:returns: Precise mimetype, or None when it is not a readable
OOXML/ODF archive.
:rtype: str
"""
try:
archive = zipfile.ZipFile(io.BytesIO(buffer))
names = archive.namelist()
except (zipfile.BadZipFile, OSError):
return None
if 'mimetype' in names:
odf = archive.read('mimetype').decode('ascii', 'ignore').strip()
if odf.startswith('application/vnd.oasis.opendocument'):
return odf
if '[Content_Types].xml' in names:
top_dirs = {name.split('/', 1)[0] for name in names}
for top_dir, mime in _OOXML_MIME_BY_TOPDIR.items():
if top_dir in top_dirs:
return mime
return None
[docs]
def magic_mime_from_buffer(buffer: bytes) -> str:
"""Detect mimetype from a file's bytes using the ``magic`` library.
OOXML (docx/xlsx/pptx) and ODF documents that libmagic reports as the
generic ``application/zip`` are refined to their precise type by
inspecting the archive -- pass the whole file, not just a header, for
that refinement to apply.
:param bytes buffer: File bytes (whole file recommended).
:returns: The detected mimetype.
:rtype: str
"""
import magic
mime = magic.from_buffer(buffer, mime=True)
if mime in _ZIP_CONTAINER_MIMES:
return _refine_zip_mime(buffer) or mime
return mime
[docs]
def guess_extension(mimetype: str) -> str:
"""Guess file extension for a mimetype.
:param str mimetype: The mimetype to look up.
:returns: File extension (including dot) or None.
:rtype: str
Example::
>>> guess_extension('image/jpeg')
'.jpg'
"""
return mimetypes.guess_extension(mimetype)
[docs]
def guess_type(url: str):
"""Guess mimetype from a URL or filename.
:param str url: URL or filename to examine.
:returns: Guessed mimetype or None.
:rtype: str
Example::
>>> guess_type('document.pdf')
'application/pdf'
>>> guess_type('image.jpg')
'image/jpeg'
"""
return mimetypes.guess_type(url)[0]
__all__ = [
'guess_type', 'guess_extension', 'magic_mime_from_buffer',
'sniff_format']
if __name__ == '__main__':
# universal type
assert guess_extension('image/jpeg') == '.jpg'
assert guess_type('a.jpg') == 'image/jpeg'
# custom patched type
assert guess_extension('x-epoc/x-sisx-app') == '.sisx'
assert guess_type('a.sisx') == 'x-epoc/x-sisx-app'