apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/io/default_data_loaders.py

@@ -0,0 +1,224 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License 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.
+
+import dataclasses
+import io
+import json
+import os
+import pathlib
+import pickle
+from collections.abc import Collection
+from typing import Any
+
+from hamilton.io.data_adapters import DataLoader, DataSaver
+from hamilton.io.utils import get_file_metadata
+
+
+@dataclasses.dataclass
+class JSONDataLoader(DataLoader):
+    path: str
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [dict, list]
+
+    def load_data(self, type_: type) -> tuple[dict, dict[str, Any]]:
+        with open(self.path, "r") as f:
+            return json.load(f), get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "json"
+
+
+@dataclasses.dataclass
+class JSONDataSaver(DataSaver):
+    path: str
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [dict, list]
+
+    @classmethod
+    def name(cls) -> str:
+        return "json"
+
+    def save_data(self, data: Any) -> dict[str, Any]:
+        with open(self.path, "w") as f:
+            json.dump(data, f)
+        return get_file_metadata(self.path)
+
+
+@dataclasses.dataclass
+class RawFileDataLoader(DataLoader):
+    path: str
+    encoding: str = "utf-8"
+
+    def load_data(self, type_: type) -> tuple[str, dict[str, Any]]:
+        with open(self.path, "r", encoding=self.encoding) as f:
+            return f.read(), get_file_metadata(self.path)
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [str]
+
+    @classmethod
+    def name(cls) -> str:
+        return "file"
+
+
+@dataclasses.dataclass
+class RawFileDataSaver(DataSaver):
+    path: str
+    encoding: str = "utf-8"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [str]
+
+    @classmethod
+    def name(cls) -> str:
+        return "file"
+
+    def save_data(self, data: Any) -> dict[str, Any]:
+        with open(self.path, "w", encoding=self.encoding) as f:
+            f.write(data)
+        return get_file_metadata(self.path)
+
+
+@dataclasses.dataclass
+class RawFileDataSaverBytes(DataSaver):
+    path: pathlib.Path | str
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [bytes, io.BytesIO]
+
+    @classmethod
+    def name(cls) -> str:
+        return "file"
+
+    def save_data(self, data: bytes | io.BytesIO) -> dict[str, Any]:
+        if isinstance(data, io.BytesIO):
+            data_bytes = data.getvalue()  # Extract bytes from BytesIO
+        else:
+            data_bytes = data
+
+        with open(self.path, "wb") as file:
+            file.write(data_bytes)
+
+        return get_file_metadata(str(self.path))
+
+
+@dataclasses.dataclass
+class PickleLoader(DataLoader):
+    path: str
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [object, Any]
+
+    @classmethod
+    def name(cls) -> str:
+        return "pickle"
+
+    def load_data(self, type_: type[object]) -> tuple[object, dict[str, Any]]:
+        with open(self.path, "rb") as f:
+            return pickle.load(f), get_file_metadata(self.path)
+
+
+@dataclasses.dataclass
+class PickleSaver(DataSaver):
+    path: str
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [object]
+
+    @classmethod
+    def name(cls) -> str:
+        return "pickle"
+
+    def save_data(self, data: Any) -> dict[str, Any]:
+        with open(self.path, "wb") as f:
+            pickle.dump(data, f)
+        return get_file_metadata(self.path)
+
+
+@dataclasses.dataclass
+class EnvVarDataLoader(DataLoader):
+    names: tuple[str, ...]
+
+    def load_data(self, type_: type[dict]) -> tuple[dict, dict[str, Any]]:
+        return {name: os.environ[name] for name in self.names}, {}
+
+    @classmethod
+    def name(cls) -> str:
+        return "environment"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [dict]
+
+
+@dataclasses.dataclass
+class LiteralValueDataLoader(DataLoader):
+    value: Any
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def load_data(self, type_: type) -> tuple[dict, dict[str, Any]]:
+        return self.value, {}
+
+    @classmethod
+    def name(cls) -> str:
+        return "literal"
+
+
+@dataclasses.dataclass
+class InMemoryResult(DataSaver):
+    """Class specifically to returning an in memory result without saving it anywhere.
+
+    Use this to get the result of a combiner easily; depends on their use.
+    """
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [Any]
+
+    def save_data(self, data: Any) -> Any:
+        return data
+
+    @classmethod
+    def name(cls) -> str:
+        return "memory"
+
+
+DATA_ADAPTERS = [
+    JSONDataSaver,
+    JSONDataLoader,
+    LiteralValueDataLoader,
+    RawFileDataLoader,
+    RawFileDataSaver,
+    RawFileDataSaverBytes,
+    PickleLoader,
+    PickleSaver,
+    EnvVarDataLoader,
+    InMemoryResult,
+]