Zenn
🌞

PythonでSpring Bootの@Profileみたいにアノテーション&環境変数だけで実装クラスをDIできるようにするには？

2022/08/14に公開
Python
Spring Boot
tech

なにこれ

表題の通り、PythonでSpring BootのようにDIするにはどうしたら良いか？の考えをまとめました。

理想像

こんなことをしたい

  • 環境変数の値に応じて、DIする実装クラスを切り替えられること
    • 切り替え先クラスを変更するたびにソースコードの修正&リリースをしたくない
    • クラウドの環境変数を変える&リスタートで反映されるようにしたい
  • 抽象クラスを継承した実装クラスにProfileアノテーションを付与するだけで、DIコンテナに登録されること
    • DIしたいクラスが増えるたびに、https://github.com/alecthomas/injector みたいにConfigurationクラスを用意して、bindする記述をしなくても動いてほしい
    • 少なくともパッケージ側で自動登録して欲しい

理想の例

Userオブジェクトを保存/取得するユーザーリポジトリクラスがあるとします。

UserRepository(抽象クラス)
import abc
from typing import NoReturn, Optional

from ... import User

class UserRepository(abc.ABC):
    @abc.abstractmethod
    def save(self, user) -> NoReturn:
        """ユーザー指定でユーザーを保存します"""
        pass

    @abc.abstractmethod
    def get(self, user_id) -> Optional[User]:
        """ユーザーID指定で該当ユーザーもしくはNoneを返却します"""
        pass
実装クラス
InMemoryUserRepository
@Profile("inMemory")
class InMemoryUserRepository(UserRepository):
    """インメモリ版のUserRepository"""
    def __init__(self):
        self.__users = {}

    @abc.abstractmethod
    def save(self, user) -> NoReturn:
        """ユーザー指定でユーザーを保存します"""
        self.__users[user.user_id] = user

    def get(self, user_id) -> Optional[User]:
        """ユーザーID指定で該当ユーザーもしくはNoneを返却します"""
        return self.__users.get(user_id, None)
PostgresqlUserRepository
@Profile("postgresql")
class PostgresqlUserRepository(UserRepository):
    """Postgresql版のUserRepository"""
    def __init__(self):
        self.__postgresql_client = # ...

    @abc.abstractmethod
    def save(self, user) -> NoReturn:
        """ユーザー指定でユーザーを保存します"""
        self.__postgresql_client.upsert(user)

    def get(self, user_id) -> Optional[User]:
        """ユーザーID指定で該当ユーザーもしくはNoneを返却します"""
        return self.__postgresql_client.find_by_id(user_id)

こんな感じで動いて欲しい！

プログラム実行時の挙動
# 環境変数が"inMemory"の時は、InMemoryUserRepositoryがDIされること
user_repository = DIContainer.get(UserRepository)
user_repository.save(User(...))　　# InMemoryUserRepository

# 環境変数が"postgresql"の時は、PostgresqlUserRepositoryがDIされること
user_repository = DIContainer.get(UserRepository)
user_repository.save(User(...))  # PostgresqlUserRepository

方法の検討

①プログラム開始時にDIコンテナに登録

  • atexitの開始時版を作る
  • プログラム開始時にProfileアノテーション側でDIコンテナに登録

atexitの開始版「atstart」を作るには、C言語が必要そう？現時点では実装難易度が高い...そもそも実現できるのか？
https://github.com/python/cpython/blob/main/Modules/atexitmodule.c

②実装クラスと検索し、DIコンテナに登録

DIContainer.getメソッド実行時にmodulefinderを使い、実装クラスを検索し、該当クラスを返す。

# この時に実装クラスを検索する
DIContainer.get(UserRepository)

モジュールやクラスが増えるたびに処理時間が長くなる...そもそも実現できるのか？

③injectorパッケージを使い、環境変数に応じて登録する実装クラスを切り替える

injectorパッケージを使い、環境変数で登録する実装クラスを切り替えられるようにする。

Injector([
    DI.new(  # Configurationクラスを生成する独自クラス
        UserRepository,   # 抽象クラス
	{
	    "inMemory": InMemoryUserRepository,   # 環境変数がinMemoryのときのDIクラス
	    "postgresql": PostgresqlUserRepository   # 環境変数がpostgresqlのときのDIクラス
	}
    ),
])

user_repository = injector.get(UserRepository)
user_repository.save(User(...))

この方法だと、環境変数に応じてDIするクラスを切り替えられる。injectorパッケージを使うので実装も用意になる。ただし、実装クラスが増えたら↑に追記する必要がある。

結論

比較項目 方法① 方法② 方法③
環境変数に応じたDIクラスの切り替えが可能 ◯ ◯ ◯
実装クラスの追加が用意 ◯ ◯ △
パフォーマンス ◯ × ◯
実現可能性 ××× △ ◯

理想としては、方法①が良いが現時点で自分にC言語の知識がないので、実現可能性は🙅
そのため、方法③が良いと思いました。

早速実装

injectorを使いながら、以下のモジュールを用意します。

.
├── di.py  # 環境変数に応じて、DIするクラスを切り替えられるようにするモジュール
└── di_manager.py  # injectorを扱うクラス
di.py
di.py
from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Type, Dict, Set, Any

from injector import Module, T, Binder, singleton


@dataclass(init=True, frozen=True)
class Profile:
    values: Set[str]

    def match(self, actives: Set[str]) -> bool:
        return all([name in actives for name in self.values])

    def __hash__(self):
        return hash(str(sorted([name for name in self.values])))

    def __eq__(self, other):
        if (other is None) or (not isinstance(other, Profile)):
            return False
        return self.values == other.values


class Switcher:
    ENV_NAME = "DI_PROFILE_ACTIVES"
    PROFILE_ACTIVES: Set[str] = set(os.getenv(ENV_NAME, "").split(","))

    @classmethod
    def get(cls, classes: Dict[Profile, Any], default: T = None) -> Any:
        for profile, a_class in classes.items():
            if profile.match(cls.PROFILE_ACTIVES):
                return a_class
        return default


@dataclass(init=True, frozen=False, unsafe_hash=True)
class DI(Module):
    interface: Type[T]
    classes: Dict[Profile, T]
    default: T = None

    @staticmethod
    def new(interface: Type[T], classes: Dict[str, T], default: T = None) -> DI:
        return DI(
            interface,
            {Profile(set(actives.split(','))): a_class for actives, a_class in classes.items()},
            default
        )

    def configure(self, binder: Binder) -> None:
        binder.bind(self.interface, to=Switcher.get(self.classes, self.default), scope=singleton)
di_manager.py
from typing import Type

from injector import Injector, T

from .di import DI
from ... import UserRepository
from ... import InMemoryUserRepository
from ... import PostgresqlUserRepository


class DIManager:
    __injector = Injector([
        # DI.new(抽象クラス, {"環境変数1", 実装クラス1, ...}),
        DI.new(UserRepository,{"inMemory": InMemoryUserRepository, "postgresql": PostgresqlUserRepository})
    ])

    @classmethod
    def get(cls, interface: Type[T]) -> T:
        return cls.__injector.get(interface)
実行時
from ... import DIManager
from ... import UserRepository

# 次の環境変数を指定
# DI_PROFILE_ACTIVES = "inMemory"

# 実装クラス = DIManager.get(抽象クラス)
user_repository = DIManager.get(UserRepository)
user_repository.save(User(...))   # InMemoryUserRepository

Discussion