Untitled

mail@pastecode.io avatar
unknown
plain_text
a month ago
1.2 kB
1
Indexable
Never
class RuleEngine:
    def __init__(self, rules):
        self.rules = rules
        self.rule_to_treatment = {
            'rule1_expr': 'treatment1_expr',
            'rule2_expr': 'treatment2_expr'
        }

    def apply_rules_and_treatments(self, df):
        lazy_df = df.lazy()
        for rule in self.rules:
            rule_methods = [method for method in dir(rule) if method.endswith('_expr')]
            for method in rule_methods:
                rule_expr = getattr(rule, method)()
                lazy_df = lazy_df.with_columns(rule_expr)
                
                # Apply treatment if rule is not respected
                if method in self.rule_to_treatment:
                    treatment_method_name = self.rule_to_treatment[method]
                    treatment_expr = getattr(rule, treatment_method_name)()
                    lazy_df = lazy_df.with_columns(treatment_expr)
        
        return lazy_df.collect()

# Example usage
rule_engine = RuleEngine(rules=[Rule1(), Rule2()])
df = pl.DataFrame({'column1': [1, 2, -1], 'column2': ['valid', 'forbidden_value', 'valid']})
result_df = rule_engine.apply_rules_and_treatments(df)
print(result_df)
Leave a Comment