def test_cycle(source): source_b64 = base64.b64encode(source) encoded_stream = io.BytesIO() with Base64IO(encoded_stream) as encoded_wrapped: encoded_wrapped.write(source) encoded_stream.seek(0) decoded_stream = io.BytesIO() with Base64IO(encoded_stream) as encoded_wrapped: decoded_stream.write(encoded_wrapped.read()) assert encoded_stream.getvalue() == source_b64 assert decoded_stream.getvalue() == source
def test_base64io_decode_with_whitespace(plaintext_source, b64_plaintext_with_whitespace, read_bytes): with Base64IO(io.BytesIO(b64_plaintext_with_whitespace)) as decoder: test = decoder.read(read_bytes) assert test == plaintext_source[:read_bytes]
def from_file( # noqa: WPS210 cls, file: Union[TextIO, BinaryIO], filename: Optional[str] = None, ) -> "File": """Convert file-like object into BotX API compatible file. Arguments: file: file-like object that will be used for creating file. filename: name that will be used for file, if was not passed, then will be retrieved from `file` `.name` property. Returns: Built file object. """ filename = filename or Path(file.name).name encoded_file = BytesIO() text_mode = file.read(0) == "" # b"" if bytes mode with Base64IO(encoded_file) as b64_stream: if text_mode: for text_line in file: # TODO: Deprecate text mode in 0.17 b64_stream.write(text_line.encode()) # type: ignore else: for line in file: b64_stream.write(line) encoded_file.seek(0) encoded_data = encoded_file.read().decode() media_type = cls._get_mimetype(filename) return cls(file_name=filename, data=cls._to_rfc2397(media_type, encoded_data))
def test_base64io_decode_parametrized_null_bytes(plaintext_source, b64_plaintext_with_whitespace, read_bytes): # Verifies that pytest is handling null bytes correctly (broken in 3.3.0) # https://github.com/pytest-dev/pytest/issues/2957 with Base64IO(io.BytesIO(b64_plaintext_with_whitespace)) as decoder: test = decoder.read(read_bytes) assert test == plaintext_source[:read_bytes]
def test_base64io_read_after_closed(): with Base64IO(io.BytesIO()) as test: with pytest.raises(ValueError) as excinfo: test.close() test.read() excinfo.match(r"I/O operation on closed file.")
async def async_from_file( # noqa: WPS210 cls, file: NamedAsyncIterable, filename: Optional[str] = None, ) -> "File": """Convert async file-like object into BotX API compatible file. Arguments: file: async file-like object that will be used for creating file. filename: name that will be used for file, if was not passed, then will be retrieved from `file` `.name` property. Returns: Built File object. """ assert hasattr( # noqa: WPS421 file, "__aiter__", ), "file should support async iteration" filename = filename or Path(file.name).name media_type = cls._get_mimetype(filename) encoded_file = BytesIO() with Base64IO(encoded_file) as b64_stream: async for line in file: # pragma: no branch b64_stream.write(line) encoded_file.seek(0) encoded_data = encoded_file.read().decode() return cls(file_name=filename, data=cls._to_rfc2397(media_type, encoded_data))
def test_base64io_encode_partial(bytes_to_generate, bytes_per_round, number_of_rounds, total_bytes_to_expect): plaintext_source = os.urandom(bytes_to_generate) plaintext_stream = io.BytesIO(plaintext_source) plaintext_b64 = base64.b64encode(plaintext_source) target_stream = io.BytesIO() target_wrapped = Base64IO(target_stream) for _round in range(number_of_rounds): target_wrapped.write(plaintext_stream.read(bytes_per_round)) # Only close if we expect to read the entire source. Otherwise, we specifically want a partial write. if bytes_to_generate == total_bytes_to_expect: target_wrapped.close() # Output length can be different if we pad, which we only do on close. expected_encoded_bytes = len(plaintext_b64) else: bytes_converted = plaintext_stream.tell() - len(target_wrapped._Base64IO__write_buffer) expected_encoded_bytes = math.ceil(bytes_converted / 3 * 4) # We read all of the bytes that we expected to from the source stream assert plaintext_stream.tell() == total_bytes_to_expect # We wrote all of the bytes that we expected to onto the target stream assert len(target_stream.getvalue()) == expected_encoded_bytes if bytes_to_generate == total_bytes_to_expect: # If we expected to process the entire stream, the results should be complete assert plaintext_b64 == target_stream.getvalue() else: # Otherwise, the encoded contents of the target stream should be a valid prefix of # the total encoded data assert plaintext_b64.startswith(target_stream.getvalue())
def test_base64io_decode_with_whitespace_str(encoding, plaintext_source, b64_plaintext_with_whitespace, read_bytes): with Base64IO(io.StringIO( b64_plaintext_with_whitespace.decode(encoding))) as decoder: test = decoder.read(read_bytes) assert test == plaintext_source[:read_bytes]
def test_base64io_decode_readline(bytes_to_read, expected_bytes_read): source_plaintext = os.urandom(io.DEFAULT_BUFFER_SIZE * 2) source_stream = io.BytesIO(base64.b64encode(source_plaintext)) with Base64IO(source_stream) as decoder: test = decoder.readline(bytes_to_read) assert test == source_plaintext[:expected_bytes_read]
def test_base64io_encode_writelines(): source_plaintext = [os.urandom(1024) for _ in range(100)] b64_plaintext = base64.b64encode(b"".join(source_plaintext)) test = io.BytesIO() with Base64IO(test) as encoder: encoder.writelines(source_plaintext) assert test.getvalue() == b64_plaintext
def test_non_interactive_error(monkeypatch, patch_method, call_method, call_arg): wrapped = io.BytesIO() monkeypatch.setattr(wrapped, patch_method, lambda: False) with Base64IO(wrapped) as wrapper: with pytest.raises(IOError) as excinfo: getattr(wrapper, call_method)(call_arg) excinfo.match(r"Stream is not " + patch_method)
def test_base64io_encode_context_manager(source_bytes): plaintext_source = os.urandom(source_bytes) plaintext_b64 = base64.b64encode(plaintext_source) plaintext_stream = io.BytesIO() with Base64IO(plaintext_stream) as plaintext_wrapped: plaintext_wrapped.write(plaintext_source) assert plaintext_stream.getvalue() == plaintext_b64
def test_base64io_decode_context_manager(): source_plaintext = os.urandom(102400) source_stream = io.BytesIO(base64.b64encode(source_plaintext)) test = io.BytesIO() with Base64IO(source_stream) as stream: for chunk in stream: test.write(chunk) assert test.getvalue() == source_plaintext
def test_base64io_decode_str(encoding, bytes_to_generate, bytes_per_round, number_of_rounds, total_bytes_to_expect): plaintext_source = os.urandom(bytes_to_generate) plaintext_b64 = io.StringIO(base64.b64encode(plaintext_source).decode(encoding)) plaintext_wrapped = Base64IO(plaintext_b64) test = b"" for _round in range(number_of_rounds): test += plaintext_wrapped.read(bytes_per_round) assert len(test) == total_bytes_to_expect assert test == plaintext_source[:total_bytes_to_expect]
def test_base64io_decode_readlines(hint_bytes, expected_bytes_read): source_plaintext = os.urandom(102400) source_stream = io.BytesIO(base64.b64encode(source_plaintext)) test = io.BytesIO() with Base64IO(source_stream) as stream: for chunk in stream.readlines(hint_bytes): test.write(chunk) assert len(test.getvalue()) == expected_bytes_read assert test.getvalue() == source_plaintext[:expected_bytes_read]
def test_base64io_encode(source_bytes): plaintext_source = os.urandom(source_bytes) plaintext_b64 = base64.b64encode(plaintext_source) plaintext_stream = io.BytesIO() plaintext_wrapped = Base64IO(plaintext_stream) try: plaintext_wrapped.write(plaintext_source) finally: plaintext_wrapped.close() assert plaintext_stream.getvalue() == plaintext_b64
def test_base64io_decode_read_only_from_buffer(): plaintext_source = b"12345" plaintext_b64 = io.BytesIO(base64.b64encode(plaintext_source)) plaintext_wrapped = Base64IO(plaintext_b64) test_1 = plaintext_wrapped.read(1) test_2 = plaintext_wrapped.read(1) test_3 = plaintext_wrapped.read() assert test_1 == b"1" assert test_2 == b"2" assert test_3 == b"345"
def test_passthrough_methods_file(tmpdir, method_name, mode, expected): source = tmpdir.join("source") source.write("some data") with open(str(source), mode) as reader: with Base64IO(reader) as b64: test = getattr(b64, method_name)() if expected: assert test else: assert not test
def _encoder(stream, should_base64): # type: (IO, bool) -> Union[IO, Base64IO] """Wraps a stream in either a Base64IO transformer or results stream if wrapping is not requested. :param stream: Stream to wrap :type stream: file-like object :param bool should_base64: Should the stream be wrapped with Base64IO :returns: wrapped stream :rtype: io.IOBase """ if should_base64: return Base64IO(stream) return stream
def test_base64io_encode_context_manager_reuse(): plaintext_source = os.urandom(10) plaintext_stream = io.BytesIO() stream = Base64IO(plaintext_stream) with stream as plaintext_wrapped: plaintext_wrapped.write(plaintext_source) with pytest.raises(ValueError) as excinfo: with stream as plaintext_wrapped: plaintext_wrapped.read() excinfo.match(r"I/O operation on closed file.")
def test_base64io_decode_file(tmpdir): source_plaintext = os.urandom(1024 * 1024) b64_plaintext = tmpdir.join("base64_plaintext") b64_plaintext.write(base64.b64encode(source_plaintext)) decoded_plaintext = tmpdir.join("decoded_plaintext") with open(str(b64_plaintext), "rb") as source: # Separate lines to accommodate 2.6 with open(str(decoded_plaintext), "wb") as raw: with Base64IO(source) as decoder: for chunk in decoder: raw.write(chunk) with open(str(decoded_plaintext), "rb") as raw: decoded = raw.read() assert decoded == source_plaintext
def wrap_stream(self, pipe: BinaryIO) -> IO: """ Wrap the given BinaryIO pipe with the appropriate stream wrapper for this method. For "RAW" or "PRINT" streams, this is a null wrapper. For BASE64 and HEX streams, this will automatically decode the data as it is streamed. Closing the wrapper will automatically close the underlying pipe. """ if self.stream is Stream.RAW or self.stream is Stream.PRINT: return pipe elif self.stream is not Stream.BASE64: raise RuntimeError( f"{self.stream.name}: we haven't implemented streaming of encodings besides base64" ) wrapped = Base64IO(pipe) original_close = wrapped.close original_write = wrapped.write def close_wrapper(): """ This is a dirty hack because Base64IO doesn't close the underlying stream when it closes. We want to assume this, so we wrap the function with one that will close the underlying stream. We need to close the Base64IO stream first, since data may be waiting to get decoded and sent. """ original_close() pipe.close() def write_wrapper(data: bytes): """ This is another nasty hack. The underlying Base64IO object erroneously returns the number of base64 bytes written, not the number if source bytes written. This causes other Python IO classes to raise an exception. We know our underlying socket will block on sending data, so all data will be sent. Again, this is gross, but it makes the python stdlib happy. """ n = original_write(data, line_length=76) return min(len(data), (n * 3) // 4) wrapped.close = close_wrapper wrapped.write = write_wrapper # We want this, but it may cause issues wrapped.name = pipe.name return wrapped
def test_base64io_encode_file(tmpdir): source_plaintext = os.urandom(1024 * 1024) plaintext_b64 = base64.b64encode(source_plaintext) plaintext = tmpdir.join("plaintext") b64_plaintext = tmpdir.join("base64_plaintext") with open(str(plaintext), "wb") as file: file.write(source_plaintext) with open(str(plaintext), "rb") as source: # Separate lines to accommodate 2.6 with open(str(b64_plaintext), "wb") as target: with Base64IO(target) as encoder: for chunk in source: encoder.write(chunk) with open(str(b64_plaintext), "rb") as file2: encoded = file2.read() assert encoded == plaintext_b64
def test_passthrough_methods_not_present(monkeypatch, method_name): wrapped = MagicMock() monkeypatch.delattr(wrapped, method_name, False) wrapper = Base64IO(wrapped) assert not getattr(wrapper, method_name)()
def test_passthrough_methods_present(monkeypatch, method_name): wrapped = io.BytesIO() monkeypatch.setattr(wrapped, method_name, lambda: sentinel.passthrough) wrapper = Base64IO(wrapped) assert getattr(wrapper, method_name)() is sentinel.passthrough
def test_unsupported_methods(method_name, args): test = Base64IO(io.BytesIO()) with pytest.raises(IOError): getattr(test, method_name)(*args)
def test_base64io_always_false_methods(method_name): test = Base64IO(io.BytesIO()) assert not getattr(test, method_name)()
def file_chunks(self) -> Generator[bytes, None, None]: """Return file data in iterator that will return bytes.""" encoded_file = BytesIO(self.data_in_base64.encode()) with Base64IO(encoded_file) as decoded_file: yield decoded_file
def test_base64io_bad_wrap(): with pytest.raises(TypeError) as excinfo: Base64IO(7) excinfo.match(r"Base64IO wrapped object must have attributes: *")
def encode_key(source_file, encoded_file): with open(source_file, "rb") as source, open(encoded_file, "wb") as target: with Base64IO(target) as encoded_target: for line in source: encoded_target.write(line)