Add Excel support and some refactoring.

Author: igorborgest

README.md

@@ -6,7 +6,7 @@
 
 > An [AWS Professional Service](https://aws.amazon.com/professional-services/) open source initiative | [email protected]
 
-[![Release](https://img.shields.io/badge/release-2.2.0-brightgreen.svg)](https://pypi.org/project/awswrangler/)
+[![Release](https://img.shields.io/badge/release-2.3.0-brightgreen.svg)](https://pypi.org/project/awswrangler/)
 [![Python Version](https://img.shields.io/badge/python-3.6%20%7C%203.7%20%7C%203.8-brightgreen.svg)](https://anaconda.org/conda-forge/awswrangler)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)

awswrangler/__metadata__.py

@@ -7,5 +7,5 @@
 
 __title__: str = "awswrangler"
 __description__: str = "Pandas on AWS."
-__version__: str = "2.2.0"
+__version__: str = "2.3.0"
 __license__: str = "Apache License 2.0"

awswrangler/s3/__init__.py

@@ -4,9 +4,11 @@
 from awswrangler.s3._delete import delete_objects  # noqa
 from awswrangler.s3._describe import describe_objects, get_bucket_region, size_objects  # noqa
 from awswrangler.s3._list import does_object_exist, list_directories, list_objects  # noqa
+from awswrangler.s3._read_excel import read_excel  # noqa
 from awswrangler.s3._read_parquet import read_parquet, read_parquet_metadata, read_parquet_table  # noqa
 from awswrangler.s3._read_text import read_csv, read_fwf, read_json  # noqa
 from awswrangler.s3._wait import wait_objects_exist, wait_objects_not_exist  # noqa
+from awswrangler.s3._write_excel import to_excel  # noqa
 from awswrangler.s3._write_parquet import store_parquet_metadata, to_parquet  # noqa
 from awswrangler.s3._write_text import to_csv, to_json  # noqa
 
@@ -32,4 +34,6 @@
     "to_parquet",
     "to_csv",
     "to_json",
+    "to_excel",
+    "read_excel",
 ]

awswrangler/s3/_fs.py

@@ -8,7 +8,7 @@
 import socket
 from contextlib import contextmanager
 from errno import ESPIPE
-from typing import Any, BinaryIO, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, cast
+from typing import Any, BinaryIO, Dict, Iterator, List, Optional, Set, Tuple, Union, cast
 
 import boto3
 from botocore.exceptions import ReadTimeoutError
@@ -188,7 +188,7 @@ def close(self) -> List[Dict[str, Union[str, int]]]:
         return self._sort_by_part_number(parts=self._results)
 
 
-class _S3ObjectBase:  # pylint: disable=too-many-instance-attributes
+class _S3ObjectBase(io.RawIOBase):  # pylint: disable=too-many-instance-attributes
     """Class to abstract S3 objects as ordinary files."""
 
     def __init__(
@@ -201,12 +201,8 @@ def __init__(
         boto3_session: Optional[boto3.Session],
         newline: Optional[str],
         encoding: Optional[str],
-        raw_buffer: bool,
     ) -> None:
-        if raw_buffer is True and "w" not in mode:
-            raise exceptions.InvalidArgumentValue("raw_buffer=True is only acceptable on write mode.")
-        self._raw_buffer: bool = raw_buffer
-        self.closed: bool = False
+        super().__init__()
         self._use_threads = use_threads
         self._newline: str = "\n" if newline is None else newline
         self._encoding: str = "utf-8" if encoding is None else encoding
@@ -408,9 +404,9 @@ def seek(self, loc: int, whence: int = 0) -> int:
 
     def flush(self, force: bool = False) -> None:
         """Write buffered data to S3."""
-        if self.closed:
+        if self.closed:  # pylint: disable=using-constant-test
             raise RuntimeError("I/O operation on closed file.")
-        if self.writable():
+        if self.writable() and self._buffer.closed is False:
             total_size: int = self._buffer.tell()
             if total_size < _MIN_WRITE_BLOCK and force is False:
                 return None
@@ -465,7 +461,7 @@ def writable(self) -> bool:
 
     def close(self) -> None:
         """Clean up the cache."""
-        if self.closed:
+        if self.closed:  # pylint: disable=using-constant-test
             return None
         if self.writable():
             _logger.debug("Closing: %s parts", self._parts_count)
@@ -488,7 +484,7 @@ def close(self) -> None:
                     ),
                 )
                 _logger.debug("complete_multipart_upload done!")
-            elif self._buffer.tell() > 0 or self._raw_buffer is True:
+            elif self._buffer.tell() > 0:
                 _logger.debug("put_object")
                 _utils.try_it(
                     f=self._client.put_object,
@@ -503,23 +499,17 @@ def close(self) -> None:
                     ),
                 )
             self._parts_count = 0
+            self._upload_proxy.close()
             self._buffer.seek(0)
             self._buffer.truncate(0)
-            self._upload_proxy.close()
             self._buffer.close()
         elif self.readable():
             self._cache = b""
         else:
             raise RuntimeError(f"Invalid mode: {self._mode}")
-        self.closed = True
+        super().close()
         return None
 
-    def get_raw_buffer(self) -> io.BytesIO:
-        """Return the Raw Buffer if it is possible."""
-        if self._raw_buffer is False:
-            raise exceptions.InvalidArgumentValue("Trying to get raw buffer with raw_buffer=False.")
-        return self._buffer
-
     def read(self, length: int = -1) -> bytes:
         """Return cached data and fetch on demand chunks."""
         if self.readable() is False:
@@ -534,8 +524,9 @@ def read(self, length: int = -1) -> bytes:
         self._loc += len(out)
         return out
 
-    def readline(self, length: int = -1) -> bytes:
+    def readline(self, length: Optional[int] = -1) -> bytes:
         """Read until the next line terminator."""
+        length = -1 if length is None else length
         end: int = self._loc + self._s3_block_size
         end = self._size if end > self._size else end
         self._fetch(self._loc, end)
@@ -553,17 +544,11 @@ def readline(self, length: int = -1) -> bytes:
             end = self._size if end > self._size else end
             self._fetch(self._loc, end)
 
-    def readlines(self) -> List[bytes]:
-        """Return all lines as list."""
-        return list(self)
-
-
-class _S3ObjectWriter(_S3ObjectBase):
-    def write(self, data: bytes) -> int:
+    def write(self, data: Union[bytes, bytearray, memoryview]) -> int:  # type: ignore
         """Write data to buffer and only upload on close() or if buffer is greater than or equal to _MIN_WRITE_BLOCK."""
         if self.writable() is False:
             raise RuntimeError("File not in write mode.")
-        if self.closed:
+        if self.closed:  # pylint: disable=using-constant-test
             raise RuntimeError("I/O operation on closed file.")
         n: int = self._buffer.write(data)
         self._loc += n
@@ -583,14 +568,12 @@ def open_s3_object(
     boto3_session: Optional[boto3.Session] = None,
     newline: Optional[str] = "\n",
     encoding: Optional[str] = "utf-8",
-    raw_buffer: bool = False,
-) -> Iterator[Union[_S3ObjectBase, _S3ObjectWriter, io.TextIOWrapper, io.BytesIO]]:
+) -> Iterator[Union[_S3ObjectBase, io.TextIOWrapper]]:
     """Return a _S3Object or TextIOWrapper based in the received mode."""
-    s3obj: Optional[Union[_S3ObjectBase, _S3ObjectWriter]] = None
+    s3obj: Optional[_S3ObjectBase] = None
     text_s3obj: Optional[io.TextIOWrapper] = None
-    s3_class: Union[Type[_S3ObjectBase], Type[_S3ObjectWriter]] = _S3ObjectWriter if "w" in mode else _S3ObjectBase
     try:
-        s3obj = s3_class(
+        s3obj = _S3ObjectBase(
             path=path,
             s3_block_size=s3_block_size,
             mode=mode,
@@ -599,11 +582,8 @@ def open_s3_object(
             boto3_session=boto3_session,
             encoding=encoding,
             newline=newline,
-            raw_buffer=raw_buffer,
         )
-        if raw_buffer is True:  # Only useful for plain io.BytesIO write
-            yield s3obj.get_raw_buffer()
-        elif "b" in mode:  # binary
+        if "b" in mode:  # binary
             yield s3obj
         else:  # text
             text_s3obj = io.TextIOWrapper(

awswrangler/s3/_read_excel.py

@@ -0,0 +1,82 @@
+"""Amazon S3 Excel Read Module (PRIVATE)."""
+
+import logging
+from typing import Any, Dict, Optional
+
+import boto3
+import pandas as pd
+
+from awswrangler import _utils, exceptions
+from awswrangler.s3._fs import open_s3_object
+
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def read_excel(
+    path: str,
+    use_threads: bool = True,
+    boto3_session: Optional[boto3.Session] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
+    **pandas_kwargs: Any,
+) -> pd.DataFrame:
+    """Read EXCEL file(s) from from a received S3 path.
+
+    Note
+    ----
+    This function accepts any Pandas's read_excel() argument.
+    https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_excel.html
+
+    Note
+    ----
+    In case of `use_threads=True` the number of threads
+    that will be spawned will be gotten from os.cpu_count().
+
+    Parameters
+    ----------
+    path : Union[str, List[str]]
+        S3 path (e.g. ``[s3://bucket/key0``).
+    use_threads : bool
+        True to enable concurrent requests, False to disable multiple threads.
+        If enabled os.cpu_count() will be used as the max number of threads.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 session will be used if boto3_session receive None.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests, only "SSECustomerAlgorithm" and "SSECustomerKey" arguments will be considered.
+    pandas_kwargs:
+        KEYWORD arguments forwarded to pandas.read_excel(). You can NOT pass `pandas_kwargs` explicit, just add valid
+        Pandas arguments in the function call and Wrangler will accept it.
+        e.g. wr.s3.read_excel(path, na_rep="", verbose=True)
+        https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_excel.html
+
+    Returns
+    -------
+    pandas.DataFrame
+        Pandas DataFrame.
+
+    Examples
+    --------
+    Reading an EXCEL file
+
+    >>> import awswrangler as wr
+    >>> df = wr.s3.read_excel('s3://bucket/key')
+
+    """
+    if "pandas_kwargs" in pandas_kwargs:
+        raise exceptions.InvalidArgument(
+            "You can NOT pass `pandas_kwargs` explicit, just add valid "
+            "Pandas arguments in the function call and Wrangler will accept it."
+            "e.g. wr.s3.read_excel(path, na_rep="
+            ", verbose=True)"
+        )
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    with open_s3_object(
+        path=path,
+        mode="rb",
+        use_threads=use_threads,
+        s3_block_size=-1,  # One shot download
+        s3_additional_kwargs=s3_additional_kwargs,
+        boto3_session=session,
+    ) as f:
+        pandas_kwargs["engine"] = "openpyxl"
+        _logger.debug("pandas_kwargs: %s", pandas_kwargs)
+        return pd.read_excel(f, **pandas_kwargs)

awswrangler/s3/_write_excel.py

@@ -0,0 +1,88 @@
+"""Amazon S3 Excel Write Module (PRIVATE)."""
+
+import logging
+from typing import Any, Dict, Optional
+
+import boto3
+import pandas as pd
+
+from awswrangler import _utils, exceptions
+from awswrangler.s3._fs import open_s3_object
+
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def to_excel(
+    df: pd.DataFrame,
+    path: str,
+    boto3_session: Optional[boto3.Session] = None,
+    s3_additional_kwargs: Optional[Dict[str, Any]] = None,
+    use_threads: bool = True,
+    **pandas_kwargs: Any,
+) -> str:
+    """Write EXCEL file on Amazon S3.
+
+    Note
+    ----
+    This function accepts any Pandas's read_excel() argument.
+    https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_excel.html
+
+    Note
+    ----
+    In case of `use_threads=True` the number of threads
+    that will be spawned will be gotten from os.cpu_count().
+
+    Parameters
+    ----------
+    df: pandas.DataFrame
+        Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html
+    path : str
+        Amazon S3 path (e.g. s3://bucket/filename.xlsx).
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 Session will be used if boto3_session receive None.
+    s3_additional_kwargs : Optional[Dict[str, Any]]
+        Forward to botocore requests. Valid parameters: "ACL", "Metadata", "ServerSideEncryption", "StorageClass",
+        "SSECustomerAlgorithm", "SSECustomerKey", "SSEKMSKeyId", "SSEKMSEncryptionContext", "Tagging".
+        e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMS_KEY_ARN'}
+    use_threads : bool
+        True to enable concurrent requests, False to disable multiple threads.
+        If enabled os.cpu_count() will be used as the max number of threads.
+    pandas_kwargs:
+        KEYWORD arguments forwarded to pandas.DataFrame.to_excel(). You can NOT pass `pandas_kwargs` explicit, just add
+        valid Pandas arguments in the function call and Wrangler will accept it.
+        e.g. wr.s3.to_excel(df, path, na_rep="", index=False)
+        https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_excel.html
+
+    Returns
+    -------
+    str
+        Written S3 path.
+
+    Examples
+    --------
+    Writing EXCEL file
+
+    >>> import awswrangler as wr
+    >>> import pandas as pd
+    >>> wr.s3.to_excel(df, 's3://bucket/filename.xlsx')
+
+    """
+    if "pandas_kwargs" in pandas_kwargs:
+        raise exceptions.InvalidArgument(
+            "You can NOT pass `pandas_kwargs` explicit, just add valid "
+            "Pandas arguments in the function call and Wrangler will accept it."
+            "e.g. wr.s3.to_excel(df, path, na_rep="
+            ", index=False)"
+        )
+    session: boto3.Session = _utils.ensure_session(session=boto3_session)
+    with open_s3_object(
+        path=path,
+        mode="wb",
+        use_threads=use_threads,
+        s3_additional_kwargs=s3_additional_kwargs,
+        boto3_session=session,
+    ) as f:
+        pandas_kwargs["engine"] = "openpyxl"
+        _logger.debug("pandas_kwargs: %s", pandas_kwargs)
+        df.to_excel(f, **pandas_kwargs)
+    return path