import json
from typing import Any, Optional

import dagster._check as check
from dagster import ConfigurableIOManager, InitResourceContext, InputContext, OutputContext, UPathIOManager
from pydantic import Field, PrivateAttr
from upath import UPath


class JSONIOManager(ConfigurableIOManager, UPathIOManager):
    base_dir: Optional[str] = Field(default=None, description="Base directory for storing files.")

    _base_path: UPath = PrivateAttr()

    def setup_for_execution(self, context: InitResourceContext) -> None:
        self._base_path = (
            UPath(self.base_dir)
            if self.base_dir is not None
            else UPath(check.not_none(context.instance).storage_directory())
        )

    def load_from_path(self, context: InputContext, path: UPath) -> str:
        with path.open("rb") as file:
            return json.loads(file)

    def dump_to_path(self, context: OutputContext, obj: Any, path: UPath):
        with path.open("wb") as file:
            json.dumps(obj, file)

UPathIOManager

•

s3:

and

s3a:

AWS S3 (requires

s3fs

to be installed)