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)
