data_utils/conf.py

from dataclasses import dataclass
from typing import Optional
from dataclasses_json import dataclass_json
import json
from pathlib import Path


@dataclass
class Slack:
    url: str
    channel: str
    username: str
    icon_emoji: str
    channel_id: Optional[str] = None
    bot_token: Optional[str] = None


class Config(object):
    def __init__(self, env="dev"):
        if env == "dev":
            config_path = Path(__file__).parent.parent / "config/config.dev.json"
        else:
            config_path = Path(__file__).parent.parent / "config/config.prod.json"
        config = self._read_config(config_path)
        self.awss3_bucket, self.awss3_path = self.__parse_awss3_storage(
            config.get("storage", {}).get("awss3", {})
        )

        (
            self.requested_img_path,
            self.text_removed_img_path,
        ) = self.__parse_local_storage(config.get("storage", {}).get("local", {}))

        self.mq_url = self._parse_amqp_server(config["mq_server"])

        (
            self.req_queue,
            self.req_pattern,
            self.resp_queue,
            self.success_resp_pattern,
            self.failure_resp_pattern,
        ) = self._parse_queue(config["queue"])

        self.slack = self._parse_slack(config["alarm"]["slack"])

    def _read_config(self, config_path) -> dict:
        with open(config_path, mode="r") as f:
            config = json.load(f)
        return config

    def _parse_amqp_server(self, amqp_server) -> str:
        username = amqp_server["username"]
        password = amqp_server["password"]
        url = amqp_server["url"]
        port = amqp_server["port"]
        amqp_url = f"amqps://{username}:{password}@{url}:{port}"
        return amqp_url

    def _parse_queue(self, queue) -> tuple:
        req_queue = queue["request_name"]
        req_pattern = queue["request_pattern"]
        resp_queue = queue.get("response_name")
        success_resp_pattern = queue["success_response_pattern"]
        failure_resp_pattern = queue["failure_response_pattern"]
        return (
            req_queue,
            req_pattern,
            resp_queue,
            success_resp_pattern,
            failure_resp_pattern,
        )

    def __parse_awss3_storage(self, awss3_storage) -> tuple:
        awss3_bucket = awss3_storage.get("default_bucket")
        awss3_path = awss3_storage.get("default_path")

        return awss3_bucket, awss3_path

    def __parse_local_storage(self, local_storage) -> tuple:
        requested_img_path = local_storage.get("requested")
        text_removed_img_path = local_storage.get("text_removed")

        return requested_img_path, text_removed_img_path

    def _parse_slack(self, slack) -> Slack:
        url = slack["url"]
        channel = slack["channel"]
        username = slack["username"]
        icon_emoji = slack["icon_emoji"]
        channel_id = slack.get("channel_id", None)
        bot_token = slack.get("bot_token", None)

        return Slack(url, channel, username, icon_emoji, channel_id, bot_token)


class ImageTrConfig(object):
    def __init__(self, env="dev"):
        if env == "dev":
            config_path = Path(__file__).parent.parent / "config/config.dev.json"
        else:
            config_path = Path(__file__).parent.parent / "config/config.prod.json"
        config = self._read_config(config_path)

        (
            self.awss3_bucket,
            self.awss3_inpainting_path,
            self.awss3_translation_path,
        ) = self.__parse_awss3_storage(config.get("storage", {}).get("awss3", {}))

        self.mq_url = self._parse_amqp_server(config["mq_server"])

        (
            self.req_queue,
            self.req_pattern,
            self.resp_queue,
            self.success_resp_pattern,
            self.failure_resp_pattern,
        ) = self._parse_queue(config["queue"])

        self.slack = self._parse_slack(config["alarm"]["slack"])

    def _read_config(self, config_path) -> dict:
        with open(config_path, mode="r") as f:
            config = json.load(f)
        return config

    def _parse_amqp_server(self, amqp_server) -> str:
        username = amqp_server["username"]
        password = amqp_server["password"]
        url = amqp_server["url"]
        port = amqp_server["port"]
        amqp_url = f"amqps://{username}:{password}@{url}:{port}"
        return amqp_url

    def _parse_queue(self, queue) -> tuple:
        req_queue = queue["request_name"]
        req_pattern = queue["request_pattern"]
        resp_queue = queue.get("response_name")
        success_resp_pattern = queue["success_response_pattern"]
        failure_resp_pattern = queue["failure_response_pattern"]
        return (
            req_queue,
            req_pattern,
            resp_queue,
            success_resp_pattern,
            failure_resp_pattern,
        )

    def __parse_awss3_storage(self, awss3_storage) -> tuple:
        awss3_bucket = awss3_storage.get("default_bucket")
        awss3_inpainting_path = awss3_storage.get("inpainting_path")
        awss3_translation_path = awss3_storage.get("translation_path")

        return awss3_bucket, awss3_inpainting_path, awss3_translation_path

    def __parse_local_storage(self, local_storage) -> tuple:
        requested_img_path = local_storage.get("requested")
        text_removed_img_path = local_storage.get("text_removed")

        return requested_img_path, text_removed_img_path

    def _parse_slack(self, slack) -> Slack:
        url = slack["url"]
        channel = slack["channel"]
        username = slack["username"]
        icon_emoji = slack["icon_emoji"]
        channel_id = slack.get("channel_id", None)
        bot_token = slack.get("bot_token", None)

        return Slack(url, channel, username, icon_emoji, channel_id, bot_token)