Stonks/overview/brokerages/registry.py

35 lines
1.1 KiB
Python

from django.conf import settings
from django.utils.module_loading import import_string
from overview.models import BrokerageConnection
from .base import BaseBrokerageAdapter, BrokerageAdapterError
from .random import RandomBrokerageAdapter
from .rest import GenericRestBrokerageAdapter
_DEFAULT_ADAPTERS = {
GenericRestBrokerageAdapter.provider: GenericRestBrokerageAdapter,
RandomBrokerageAdapter.provider: RandomBrokerageAdapter,
}
def register_brokerage_adapter(provider: str, adapter_class: type[BaseBrokerageAdapter]):
_DEFAULT_ADAPTERS[provider] = adapter_class
def get_brokerage_adapter(connection: BrokerageConnection) -> BaseBrokerageAdapter:
adapter_class = _DEFAULT_ADAPTERS.get(connection.provider)
configured_paths = getattr(settings, "BROKERAGE_ADAPTERS", {})
configured_path = configured_paths.get(connection.provider)
if configured_path:
adapter_class = import_string(configured_path)
if not adapter_class:
raise BrokerageAdapterError(
f"No brokerage adapter is registered for provider '{connection.provider}'."
)
return adapter_class(connection)