Fix test related problems and enhance zstd support

filetype/types/archive.py

@@ -2,6 +2,8 @@
 
 from __future__ import absolute_import
 
+import struct
+
 from .base import Type
 
 
@@ -188,7 +190,7 @@ def match(self, buf):
         if (len(buf) > 3 and
             buf[0] == 0xEF and
             buf[1] == 0xBB and
-            buf[2] == 0xBF):
+            buf[2] == 0xBF):  # noqa E129
             buf = buf[3:]
 
         return (len(buf) > 3 and
@@ -644,16 +646,42 @@ class Zstd(Type):
     """
     MIME = 'application/zstd'
     EXTENSION = 'zst'
+    MAGIC_SKIPPABLE_START = 0x184D2A50
+    MAGIC_SKIPPABLE_MASK = 0xFFFFFFF0
 
     def __init__(self):
         super(Zstd, self).__init__(
             mime=Zstd.MIME,
             extension=Zstd.EXTENSION
         )
 
+    @staticmethod
+    def _to_little_endian_int(buf):
+        # return int.from_bytes(buf, byteorder='little')
+        return struct.unpack('<L', buf)[0]
+
     def match(self, buf):
-        return (len(buf) > 3 and
-                buf[0] in (0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28) and
-                buf[1] == 0xb5 and
-                buf[2] == 0x2f and
-                buf[3] == 0xfd)
+        # Zstandard compressed data is made of one or more frames.
+        # There are two frame formats defined by Zstandard:
+        # Zstandard frames and Skippable frames.
+        # See more details from
+        # https://tools.ietf.org/id/draft-kucherawy-dispatch-zstd-00.html#rfc.section.2
+        is_zstd = (
+            len(buf) > 3 and
+            buf[0] in (0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28) and
+            buf[1] == 0xb5 and
+            buf[2] == 0x2f and
+            buf[3] == 0xfd)
+        if is_zstd:
+            return True
+        # skippable frames
+        if len(buf) < 8:
+            return False
+        magic = self._to_little_endian_int(buf[:4]) & Zstd.MAGIC_SKIPPABLE_MASK
+        if magic == Zstd.MAGIC_SKIPPABLE_START:
+            user_data_len = self._to_little_endian_int(buf[4:8])
+            if len(buf) < 8 + user_data_len:
+                return False
+            next_frame = buf[8 + user_data_len:]
+            return self.match(next_frame)
+        return False

requirements-dev.txt

@@ -1,2 +1,3 @@
 pytest
+pytest-benchmark
 flake8

tests/test_types.py

@@ -61,10 +61,11 @@ def test_guess_mov(self):
         self.assertEqual(kind.extension, 'mov')
 
     def test_guess_zstd(self):
-        kind = filetype.guess(FIXTURES + '/sample.zst')
-        self.assertTrue(kind is not None)
-        self.assertEqual(kind.mime, 'application/zstd')
-        self.assertEqual(kind.extension, 'zst')
+        for name in 'sample.zst', 'sample_skippable.zst':
+            kind = filetype.guess(FIXTURES + '/' + name)
+            self.assertTrue(kind is not None)
+            self.assertEqual(kind.mime, 'application/zstd')
+            self.assertEqual(kind.extension, 'zst')
 
     def test_guess_doc(self):
         kind = filetype.guess(FIXTURES + '/sample.doc')

tox.ini

@@ -7,23 +7,24 @@
 envlist = py{27,35,36,37,38,39}, lint, doc, clean
 skip_missing_interpreters = true
 
-[testenv]
-commands = python -m unittest discover
+[testenv:test]
+deps = pytest-benchmark
+commands = pytest
 
 [testenv:lint]
 basepython = python3
 deps =
      flake8
-commands = flake8 {toxinidir}
+commands = flake8 {toxinidir} --extend-exclude tests,docs,build,dist,venv
 
 [testenv:doc]
 basepython = python3
 deps =
-     pdoc
-commands = pdoc --html --overwrite --all-submodules --html-dir docs filetype
+     pdoc3
+commands = pdoc --html --force --output-dir docs filetype
 
 [testenv:clean]
 basepython = python3
 deps =
      pyclean
-commands = py3clean {toxinidir}
+commands = pyclean {toxinidir}