diff --git a/dol/util.py b/dol/util.py index 77d2978b..0d7d4af7 100644 --- a/dol/util.py +++ b/dol/util.py @@ -1697,7 +1697,9 @@ def written_bytes( and returns a function that instead of writing to a file, returns the bytes that would have been written. - Note: If obj is not given, write_bytes will return a "bytes writer" function that + This is the write version of the `read_from_bytes` function of the same module. + + Note: If obj is not given, `write_bytes` will return a "bytes writer" function that takes obj as the first argument, and uses the file_writer to write the bytes. :param file_writer: A function that writes an object to a file-like object. @@ -1756,6 +1758,98 @@ def written_bytes( return buffer.read() +def _call_reader( + reader: Callable, + buffer: Buffer, + buffer_arg_position: int = 0, + buffer_arg_name: str = None, + *args, + **kwargs, +): + """ + Helper function to handle reading from the buffer based on buffer_arg_position or buffer_arg_name. + + :param reader: A function that reads from a file-like object. + :param buffer: The file-like object to read from. + :param buffer_arg_position: Position of the file-like object argument in reader function. + :param buffer_arg_name: Name of the file-like object argument in reader function. + :raises ValueError: If buffer_arg_position is not valid. + """ + if buffer_arg_name is not None: + kwargs[buffer_arg_name] = buffer + return reader(*args, **kwargs) + else: + args = list(args) + # Ensure the args list is long enough + while len(args) < buffer_arg_position: + args.append(None) + args.insert(buffer_arg_position, buffer) + return reader(*args, **kwargs) + + +def read_from_bytes( + file_reader: Callable, + obj: bytes = None, + *, + buffer_arg_position: int = 0, + buffer_arg_name: str = None, + io_buffer_cls: Buffer = io.BytesIO, + **kwargs, +): + """ + Takes a file reading function that expects a file-like object, + and returns a function that instead of reading from a file, reads from bytes. + + This is the read version of the `written_bytes` function of the same module. + + Note: If obj is not given, read_from_bytes will return a "bytes reader" function that + takes obj as the first argument, and uses the file_reader to read the bytes. + + :param file_reader: A function that reads from a file-like object. + :param obj: The bytes to read. + :param buffer_arg_position: The position of the file-like object in file_reader's arguments. + :param buffer_arg_name: The name of the file-like object argument in file_reader. + :return: The result of reading from the bytes. + + Example usage: + + Using `json.load` to read a JSON object from bytes: + + >>> import json + >>> data = {'a': 1, 'b': 2} + >>> json_bytes = json.dumps(data).encode('utf-8') + >>> read_json_from_bytes = read_from_bytes(json.load) + >>> data_loaded = read_json_from_bytes(json_bytes) + >>> data_loaded == data + True + + Using `pickle.load` to read an object from bytes: + + >>> import pickle + >>> obj = {'x': [1, 2, 3], 'y': ('a', 'b')} + >>> pickle_bytes = pickle.dumps(obj) + >>> read_pickle_from_bytes = read_from_bytes(pickle.load) + >>> obj_loaded = read_pickle_from_bytes(pickle_bytes) + >>> obj_loaded == obj + True + """ + if obj is None: + return partial( + read_from_bytes, + file_reader, + buffer_arg_position=buffer_arg_position, + buffer_arg_name=buffer_arg_name, + io_buffer_cls=io_buffer_cls, + **kwargs, + ) + + buffer = io_buffer_cls(obj) + + return _call_reader( + file_reader, buffer, buffer_arg_position, buffer_arg_name, **kwargs + ) + + def write_to_file(obj: VT, key: KT): with open(key, "wb") as f: f.write(obj)