easy-ml/feature/feature_strategy_factory.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/11/25
@desc: 特征筛选策略工厂
"""
from typing import Type

from commom import GeneralException
from enums import FeatureStrategyEnum, ResultCodesEnum
from .bin.strategy_norm import StrategyNorm
from .feature_strategy_base import FeatureStrategyBase
from .woe.strategy_woe import StrategyWoe

strategy_map = {
    FeatureStrategyEnum.WOE.value: StrategyWoe,
    FeatureStrategyEnum.NORM.value: StrategyNorm
}


class FeatureStrategyFactory():

    @staticmethod
    def get_strategy(strategy: str) -> Type[FeatureStrategyBase]:
        if strategy not in strategy_map.keys():
            raise GeneralException(ResultCodesEnum.NOT_FOUND, message=f"特征搜索策略【{strategy}】不存在")
        strategy = strategy_map.get(strategy)
        return strategy