models.py

from typing import Dict, List, override
import pydantic
import enum


class ViolationKeyType(enum.Enum):
    OPERATION_ID = "operation_id"
    PATH = "path"
    SCHEMA_PATH = "schema_path"
    HTTP_CODE = "http_code"


class ViolationKey(pydantic.BaseModel, frozen=True):
    # One of:
    # (for operation-level rules)
    # - operation_id: str
    # or
    # (for operation-level rules)
    # - path: str
    # - method: str
    # or
    # (for path-level rules)
    # - path: str
    # or
    # (for schema rules)
    # - schema_path: str
    # or
    # (for http code rules)
    # - http_code: str

    violation_key_type: ViolationKeyType
    operation_id: str | None = None
    path: str | None = None
    schema_path: str | None = None
    http_code: str | None = None
    
    def equal(self, other: "ViolationKey") -> bool:
        """
        Check if this ViolationKey is equal to or a subset of another ViolationKey.
        """
        if self == other:
            return True
        if self.violation_key_type != other.violation_key_type:
            return False
        if self.violation_key_type == ViolationKeyType.OPERATION_ID:
            return self.operation_id == other.operation_id
        if self.violation_key_type == ViolationKeyType.PATH:
            return self.path == other.path
        if self.violation_key_type == ViolationKeyType.SCHEMA_PATH:
            return self.schema_path == other.schema_path
        if self.violation_key_type == ViolationKeyType.HTTP_CODE:
            return self.http_code == other.http_code
        return False

    @pydantic.model_validator(mode="after")
    def validate_violation_key(self) -> "ViolationKey":
        required_fields = {
            ViolationKeyType.OPERATION_ID: ["operation_id"],
            ViolationKeyType.PATH: ["path"],
            ViolationKeyType.SCHEMA_PATH: ["schema_path"],
            ViolationKeyType.HTTP_CODE: ["http_code"],
        }

        fields = required_fields[self.violation_key_type]
        missing = [field for field in fields if getattr(self, field) is None]
        if missing:
            raise ValueError(f"Missing required fields for {self.violation_key_type.value}: {missing}")

        allowed = set(fields + ["violation_key_type"])
        extras = [name for name, value in self.model_dump().items() if name not in allowed and value is not None]
        if extras:
            raise ValueError(f"Unexpected fields set for {self.violation_key_type.value}: {extras}")

        return self

    @override
    def model_dump(self, *args, **kwargs):
        # Override the model_dump method to exclude None values
        kwargs.pop("exclude_none", None)
        return super().model_dump(*args, exclude_none=True, **kwargs)


class RuleSource(pydantic.BaseModel):
    name: str


class Violation(pydantic.BaseModel):
    rule_id: int
    message: str
    source: RuleSource | None = None


class RulesConfig(pydantic.BaseModel):
    config: Dict[int, Dict[str, str]] = {}
    enabled: Dict[int, bool] = {}
    ignore: Dict[int, List[ViolationKey]] = {}
    ignore_all: List[ViolationKey] = []

    def disable_rule(self, rule_id: int) -> None:
        self.enabled[rule_id] = False

    def enable_rule(self, rule_id: int) -> None:
        self.enabled[rule_id] = True

class RuleMetadata(pydantic.BaseModel):
    rule_id: int
    name: str
    description: str
    hidden: bool

class ExtendedRulesConfig(RulesConfig):
    available_config: Dict[int, Dict[str, List[str | None]]]
    rules_metadata: Dict[int, RuleMetadata]


    @classmethod
    def from_config(
        cls,
        config: RulesConfig,
        available_config: Dict[int, Dict[str, List[str | None]]],
        rules_metadata: Dict[int, RuleMetadata],
    ) -> "ExtendedRulesConfig":
        return cls(
            config=config.config,
            enabled=config.enabled,
            ignore=config.ignore,
            ignore_all=config.ignore_all,
            available_config=available_config,
            rules_metadata=rules_metadata,
        )