pickle.py 1.1 KB

1234567891011121314151617181920212223242526272829303132
  1. import pickle
  2. from typing import Any
  3. from django.core.exceptions import ImproperlyConfigured
  4. from .base import BaseSerializer
  5. class PickleSerializer(BaseSerializer):
  6. def __init__(self, options) -> None:
  7. self._pickle_version = pickle.DEFAULT_PROTOCOL
  8. self.setup_pickle_version(options)
  9. super().__init__(options=options)
  10. def setup_pickle_version(self, options) -> None:
  11. if "PICKLE_VERSION" in options:
  12. try:
  13. self._pickle_version = int(options["PICKLE_VERSION"])
  14. if self._pickle_version > pickle.HIGHEST_PROTOCOL:
  15. raise ImproperlyConfigured(
  16. f"PICKLE_VERSION can't be higher than pickle.HIGHEST_PROTOCOL:"
  17. f" {pickle.HIGHEST_PROTOCOL}"
  18. )
  19. except (ValueError, TypeError):
  20. raise ImproperlyConfigured("PICKLE_VERSION value must be an integer")
  21. def dumps(self, value: Any) -> bytes:
  22. return pickle.dumps(value, self._pickle_version)
  23. def loads(self, value: bytes) -> Any:
  24. return pickle.loads(value)