| import json | |
| import pickle | |
| from enum import Enum | |
| from pathlib import Path | |
| from typing import Any, Callable, Union | |
| from pydantic.v1.types import StrBytes | |
| class Protocol(str, Enum): | |
| json = 'json' | |
| pickle = 'pickle' | |
| def load_str_bytes( | |
| b: StrBytes, | |
| *, | |
| content_type: str = None, | |
| encoding: str = 'utf8', | |
| proto: Protocol = None, | |
| allow_pickle: bool = False, | |
| json_loads: Callable[[str], Any] = json.loads, | |
| ) -> Any: | |
| if proto is None and content_type: | |
| if content_type.endswith(('json', 'javascript')): | |
| pass | |
| elif allow_pickle and content_type.endswith('pickle'): | |
| proto = Protocol.pickle | |
| else: | |
| raise TypeError(f'Unknown content-type: {content_type}') | |
| proto = proto or Protocol.json | |
| if proto == Protocol.json: | |
| if isinstance(b, bytes): | |
| b = b.decode(encoding) | |
| return json_loads(b) | |
| elif proto == Protocol.pickle: | |
| if not allow_pickle: | |
| raise RuntimeError('Trying to decode with pickle with allow_pickle=False') | |
| bb = b if isinstance(b, bytes) else b.encode() | |
| return pickle.loads(bb) | |
| else: | |
| raise TypeError(f'Unknown protocol: {proto}') | |
| def load_file( | |
| path: Union[str, Path], | |
| *, | |
| content_type: str = None, | |
| encoding: str = 'utf8', | |
| proto: Protocol = None, | |
| allow_pickle: bool = False, | |
| json_loads: Callable[[str], Any] = json.loads, | |
| ) -> Any: | |
| path = Path(path) | |
| b = path.read_bytes() | |
| if content_type is None: | |
| if path.suffix in ('.js', '.json'): | |
| proto = Protocol.json | |
| elif path.suffix == '.pkl': | |
| proto = Protocol.pickle | |
| return load_str_bytes( | |
| b, proto=proto, content_type=content_type, encoding=encoding, allow_pickle=allow_pickle, json_loads=json_loads | |
| ) | |