diff --git a/pyathena/filesystem/s3_object.py b/pyathena/filesystem/s3_object.py index dc6ffbdfc..825f2aac3 100644 --- a/pyathena/filesystem/s3_object.py +++ b/pyathena/filesystem/s3_object.py @@ -33,6 +33,19 @@ "Metadata": "metadata", "LastModified": "last_modified", } +# Fields read as None through attribute access when the object does not have them. +_S3_OBJECT_FIELDS = frozenset( + [ + *_API_FIELD_TO_S3_OBJECT_PROPERTY.values(), + "name", + "type", + "bucket", + "key", + "size", + "version_id", + "is_latest", + ] +) class S3ObjectType: @@ -95,7 +108,10 @@ class S3Object(MutableMapping[str, Any]): The object supports both dictionary-style access and property-style access to metadata fields like content type, storage class, encryption - settings, and object lock configurations. + settings, and object lock configurations. Dictionary-style access + behaves like a dictionary, so a missing key raises KeyError. Property-style + access returns None for a known field that the object does not have, + and raises AttributeError for any other missing name. Example: >>> s3_obj = S3Object({"ContentType": "text/csv", "ContentLength": 1024}) @@ -163,10 +179,26 @@ def get(self, key: str, default: Any = None) -> Any: @override def __getitem__(self, item: str) -> Any: - return self.__dict__.get(item) + return self.__dict__[item] - def __getattr__(self, item: str): - return self.get(item) + def __getattr__(self, item: str) -> Any: + """Return None for a known field that the object does not have. + + Called only when normal attribute lookup fails, so fields that the + object has are returned without reaching this method. + + Args: + item: The attribute name. + + Returns: + None, if ``item`` is a known S3 object field. + + Raises: + AttributeError: If ``item`` is not a known S3 object field. + """ + if item in _S3_OBJECT_FIELDS: + return None + raise AttributeError(f"{type(self).__name__!r} object has no attribute {item!r}") @override def __setitem__(self, key: str, value: Any) -> None: @@ -192,6 +224,14 @@ def __len__(self) -> int: def __str__(self): return str(self.__dict__) + def copy(self) -> S3Object: + """Return a shallow copy of the object. + + Returns: + A new S3Object with the same fields. + """ + return copy.copy(self) + def to_dict(self) -> dict[str, Any]: """Convert S3Object to dictionary representation. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 57d18cd92..bd238ffb3 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -22,6 +22,7 @@ import pytest from fsspec import Callback from fsspec.dircache import DirCache +from fsspec.implementations.dirfs import DirFileSystem import pyathena from pyathena.filesystem import register_s3_filesystem @@ -1446,6 +1447,33 @@ def test_ls_versions_object_path_falls_back_to_the_key(self): ("bucket/path/key", "v1", 2), ] + def test_dir_filesystem(self): + # DirFileSystem copies every entry with copy() before renaming it. + fs = self._make_fs() + fs._call.side_effect = [ + { + "CommonPrefixes": [{"Prefix": "path/dir/"}], + "Contents": [{"Key": "path/key", "Size": 4}], + "IsTruncated": False, + }, + {"ContentLength": 4, "ETag": '"etag"'}, + ] + dir_fs = DirFileSystem(path="bucket/path", fs=fs) + + actual = dir_fs.ls("", detail=True) + assert [(f["name"], f["type"]) for f in actual] == [("dir", "directory"), ("key", "file")] + assert all(isinstance(f, S3Object) for f in actual) + actual = dir_fs.info("key") + assert isinstance(actual, S3Object) + assert (actual.name, actual.size) == ("key", 4) + # The cached entries keep their full names. + assert [f.name for f in fs.ls("bucket/path", detail=True)] == [ + "bucket/path/dir", + "bucket/path/key", + ] + assert fs.info("bucket/path/key").name == "bucket/path/key" + assert fs._call.call_count == 2 + def test_metadata_with_version_id(self): fs = self._make_fs() fs._call.return_value = {"Metadata": {}} diff --git a/tests/pyathena/filesystem/test_s3_object.py b/tests/pyathena/filesystem/test_s3_object.py index 51afb2072..a6035419c 100644 --- a/tests/pyathena/filesystem/test_s3_object.py +++ b/tests/pyathena/filesystem/test_s3_object.py @@ -5,8 +5,12 @@ # # SPDX-License-Identifier: MIT +import copy +import pickle from datetime import datetime +import pytest + from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -103,6 +107,61 @@ def test_to_api_repr(self): "StorageClass": "STANDARD", } + @staticmethod + def _file_object(): + return S3Object( + init={"ContentLength": 3, "ETag": '"etag"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="test-bucket", + key="path/to/object", + ) + + def test_mapping(self): + actual = self._file_object() + assert actual["etag"] == '"etag"' + with pytest.raises(KeyError): + actual["version_id"] + assert "etag" in actual + assert "version_id" not in actual + assert actual.get("etag") == '"etag"' + assert actual.get("version_id") is None + assert actual.get("version_id", "default") == "default" + assert actual.pop("version_id", "default") == "default" + assert actual.setdefault("version_id", "v1") == "v1" + assert actual["version_id"] == "v1" + + def test_attribute(self): + actual = self._file_object() + assert actual.etag == '"etag"' + # Known fields that the object does not have read as None. + assert actual.version_id is None + assert actual.is_latest is None + assert actual.content_type is None + assert "content_type" not in actual + with pytest.raises(AttributeError): + _ = actual.unknown + assert not hasattr(actual, "__setstate__") + + @pytest.mark.parametrize( + "func", + [ + copy.copy, + copy.deepcopy, + lambda obj: pickle.loads(pickle.dumps(obj)), + lambda obj: obj.copy(), + ], + ids=["copy", "deepcopy", "pickle", "copy_method"], + ) + def test_copy(self, func): + expected = self._file_object() + actual = func(expected) + assert isinstance(actual, S3Object) + assert actual is not expected + assert actual == expected + assert actual.name == "test-bucket/path/to/object" + actual["name"] = "renamed" + assert expected.name == "test-bucket/path/to/object" + class TestS3Metadata: def test_init(self):