diff --git a/src/base64io/__init__.py b/src/base64io/__init__.py index e687913..b6a4a78 100644 --- a/src/base64io/__init__.py +++ b/src/base64io/__init__.py @@ -43,6 +43,15 @@ def _py2(): return sys.version_info[0] == 2 +if not _py2(): + # The "file" object does not exist in Python 3, but we need to reference + # it in Python 2 code paths. Defining this here accomplishes two things: + # First, it allows linters to accept "file" as a defined object in Python 3. + # Second, it will serve as a canary to ensure that there are no references + # to "file" in Python 3 code paths. + file = NotImplemented # pylint: disable=invalid-name + + class Base64IO(io.IOBase): """Base64 stream with context manager support. @@ -121,9 +130,8 @@ class Base64IO(io.IOBase): try: method = getattr(self.__wrapped, method_name) except AttributeError: - if _py2() and isinstance(self.__wrapped, file): # type: ignore # noqa pylint: disable=undefined-variable - if mode in self.__wrapped.mode: - return True + if _py2() and isinstance(self.__wrapped, file) and mode in self.__wrapped.mode: + return True return False else: return method() @@ -316,9 +324,15 @@ class Base64IO(io.IOBase): :rtype: list of bytes """ lines = [] + total_len = 0 + hint_defined = hint > 0 + for line in self: # type: ignore lines.append(line) - if hint > 0 and len(lines) * io.DEFAULT_BUFFER_SIZE > hint: + total_len += len(line) + + hint_satisfied = total_len > hint + if hint_defined and hint_satisfied: break return lines diff --git a/test/unit/test_base64_stream.py b/test/unit/test_base64_stream.py index b0f114a..de8410c 100644 --- a/test/unit/test_base64_stream.py +++ b/test/unit/test_base64_stream.py @@ -10,7 +10,7 @@ # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -"""Unit test suite for ``aws_encryption_sdk_cli.internal.encoding``.""" +"""Unit test suite for ``base64io.Base64IO``.""" from __future__ import division import base64 diff --git a/test/unit/test_base64io.py b/test/unit/test_base64io.py new file mode 100644 index 0000000..dfb9d20 --- /dev/null +++ b/test/unit/test_base64io.py @@ -0,0 +1,39 @@ +# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit test suite for ``base64io`` helpers.""" +import sys + +import pytest + +import base64io + +pytestmark = [pytest.mark.unit] + + +def test_py2(): + is_python2 = sys.version_info[0] == 2 + + assert (is_python2 and base64io._py2()) or (not is_python2 and not base64io._py2()) + + +def test_file(): + is_python2 = sys.version_info[0] == 2 + + if is_python2: + # If we are in Python 2, the "file" assignment should not + # happen because it is a builtin object. + assert not hasattr(base64io, 'file') + else: + # If we are in Python 3, the "file" assignment should happen + # to provide a concrete definition of the "file" name. + assert base64io.file is NotImplemented