Mirror of https://github.com/roostorg/osprey
github.com/roostorg/osprey
1from dataclasses import dataclass
2from typing import List, Self, cast
3
4from osprey.engine.executor.custom_extracted_features import CustomExtractedFeature
5from osprey.engine.executor.execution_context import ExecutionContext
6from osprey.engine.language_types.effects import EffectBase, EffectToCustomExtractedFeatureBase
7from osprey.engine.stdlib.udfs.categories import UdfCategories
8from osprey.engine.udf.arguments import ArgumentsBase
9from osprey.engine.udf.base import UDFBase
10from osprey.engine.utils.types import add_slots
11
12
13class BanUserArguments(ArgumentsBase):
14 entity: str
15 comment: str
16
17
18@dataclass
19class BanUserEffect(EffectToCustomExtractedFeatureBase[List[str]]):
20 """Adds a 'ban user' effect to the action."""
21
22 entity: str
23 """The entity that the ban will be applied to."""
24
25 comment: str
26 """The comment that will be included with the ban, for internal purposes."""
27
28 def to_str(self) -> str:
29 return f'{self.entity}|{self.comment}'
30
31 @classmethod
32 def build_custom_extracted_feature_from_list(cls, values: List[Self]) -> CustomExtractedFeature[List[str]]:
33 return BanEffectsExtractedFeature(effects=cast(List[BanUserEffect], values))
34
35
36@add_slots
37@dataclass
38class BanEffectsExtractedFeature(CustomExtractedFeature[List[str]]):
39 effects: List[BanUserEffect]
40
41 @classmethod
42 def feature_name(cls) -> str:
43 return 'ban_user'
44
45 def get_serializable_feature(self) -> List[str] | None:
46 return [effect.to_str() for effect in self.effects]
47
48
49def synthesize_effect(arguments: BanUserArguments) -> BanUserEffect:
50 return BanUserEffect(
51 entity=arguments.entity,
52 comment=arguments.comment,
53 )
54
55
56class BanUser(UDFBase[BanUserArguments, EffectBase]):
57 category = UdfCategories.ENGINE
58
59 def execute(self, execution_context: ExecutionContext, arguments: BanUserArguments) -> EffectBase:
60 return synthesize_effect(arguments)