diff --git a/fileformats/datascience/__init__.py b/fileformats/datascience/__init__.py index af22e17..0ebedf3 100644 --- a/fileformats/datascience/__init__.py +++ b/fileformats/datascience/__init__.py @@ -11,7 +11,7 @@ from .scripts import RFile, IPythonNotebook from .object_serialization import ( Pickle, - Pickle___Gzip, + Pickle__Gzip, ) __all__ = [ @@ -26,5 +26,5 @@ "RFile", "IPythonNotebook", "Pickle", - "Pickle___Gzip", + "Pickle__Gzip", ] diff --git a/fileformats/datascience/object_serialization.py b/fileformats/datascience/object_serialization.py index 31a742e..4355b5a 100644 --- a/fileformats/datascience/object_serialization.py +++ b/fileformats/datascience/object_serialization.py @@ -1,19 +1,26 @@ -from fileformats.generic import BinaryFile -from fileformats.core.mixin import WithMagicNumber from fileformats.application import Gzip +from fileformats.core.mixin import WithMagicVersion +from fileformats.generic import BinaryFile class ObjectSerialisation(BinaryFile): ... -class Pickle(WithMagicNumber, ObjectSerialisation): +class Pickle(WithMagicVersion, ObjectSerialisation): """Python's native byte-encoded serialization format""" ext = ".pkl" - magic_number = "8004" + magic_pattern = rb"\x80([\x02-\xff])" + magic_pattern_maxlength = 2 + + @classmethod + def decode_version(cls, version_bytes: bytes) -> str: + # The protocol version is stored as the integer value + # of the byte + return str(version_bytes[0]) -class Pickle___Gzip(Gzip[Pickle]): # type: ignore[type-arg] +class Pickle__Gzip(Gzip[Pickle]): # type: ignore[type-arg] """Python pickle file that has been gzipped""" ext = "pkl.gz" diff --git a/fileformats/datascience/tests/test_mime.py b/fileformats/datascience/tests/test_mime.py index 29aa51d..ce23a47 100644 --- a/fileformats/datascience/tests/test_mime.py +++ b/fileformats/datascience/tests/test_mime.py @@ -1,8 +1,8 @@ from fileformats.core import from_mime -from fileformats.datascience import Pickle___Gzip +from fileformats.datascience import Pickle__Gzip def test_native_container_roundtrip() -> None: - mime = Pickle___Gzip.mime_like - assert Pickle___Gzip is from_mime(mime) + mime = Pickle__Gzip.mime_like + assert Pickle__Gzip is from_mime(mime)