From 9d903a06317bea87c29b1a586dcb96e2251f51e1 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 24 May 2023 20:46:07 +0300
Subject: [PATCH] Add support for puremagic

---
 mautrix/util/magic.py     | 40 ++++++++++++++++++++++++++++++++-------
 optional-requirements.txt |  2 +-
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/mautrix/util/magic.py b/mautrix/util/magic.py
index 5c061993..38f4894f 100644
--- a/mautrix/util/magic.py
+++ b/mautrix/util/magic.py
@@ -7,14 +7,40 @@
 
 import functools
 
-import magic
-
 try:
-    _from_buffer = functools.partial(magic.from_buffer, mime=True)
-    _from_filename = functools.partial(magic.from_file, mime=True)
-except AttributeError:
-    _from_buffer = lambda data: magic.detect_from_content(data).mime_type
-    _from_filename = lambda file: magic.detect_from_filename(file).mime_type
+    from puremagic.main import PureMagicWithConfidence as _PureMagicResult
+    import puremagic
+
+    _blacklist = {".koz"}
+
+    def _cleanup_result(results: list[_PureMagicResult]) -> str:
+        for res in results:
+            if not res.mime_type or res.extension in _blacklist:
+                continue
+            return res.mime_type
+        return "application/octet-stream"
+
+    def _from_buffer(data: bytes) -> str:
+        try:
+            return _cleanup_result(puremagic.magic_string(data))
+        except puremagic.PureError:
+            return "application/octet-stream"
+
+    def _from_filename(filename: str) -> str:
+        try:
+            return _cleanup_result(puremagic.magic_file(filename))
+        except puremagic.PureError:
+            return "application/octet-stream"
+
+except ImportError:
+    import magic
+
+    try:
+        _from_buffer = functools.partial(magic.from_buffer, mime=True)
+        _from_filename = functools.partial(magic.from_file, mime=True)
+    except AttributeError:
+        _from_buffer = lambda data: magic.detect_from_content(data).mime_type
+        _from_filename = lambda file: magic.detect_from_filename(file).mime_type
 
 
 def mimetype(data: bytes | bytearray | str) -> str:
diff --git a/optional-requirements.txt b/optional-requirements.txt
index a5660f4a..67d7a3a0 100644
--- a/optional-requirements.txt
+++ b/optional-requirements.txt
@@ -1,4 +1,4 @@
-python-magic
+puremagic
 ruamel.yaml
 SQLAlchemy<2
 commonmark