|
- 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
- )
|
