diff --git a/src/dbt_score/rule.py b/src/dbt_score/rule.py index 686bb49..ea145ce 100644 --- a/src/dbt_score/rule.py +++ b/src/dbt_score/rule.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from enum import Enum -from typing import Callable, Type +from typing import Any, Callable, Type from dbt_score.models import Model @@ -43,7 +43,7 @@ def evaluate(self, model: Model) -> RuleViolation | None: def rule( description: str | None = None, severity: Severity = Severity.MEDIUM, -) -> Callable[[Callable[[Model], RuleViolation | None]], Type[Rule]]: +) -> Callable[[Callable[[Any, Model], RuleViolation | None]], Type[Rule]]: """Rule decorator. The rule decorator creates a rule class (subclass of Rule) and returns it. @@ -54,7 +54,7 @@ def rule( """ def decorator_rule( - func: Callable[[Model], RuleViolation | None], + func: Callable[[Any, Model], RuleViolation | None], ) -> Type[Rule]: """Decorator function.""" if func.__doc__ is None and description is None: diff --git a/src/dbt_score/rules/example_rules.py b/src/dbt_score/rules/example_rules.py index 333fabd..0189cff 100644 --- a/src/dbt_score/rules/example_rules.py +++ b/src/dbt_score/rules/example_rules.py @@ -24,7 +24,7 @@ def evaluate(self, model: Model) -> RuleViolation | None: @rule() -def has_owner(self, model: Model) -> RuleViolation | None: +def has_owner(self: Rule, model: Model) -> RuleViolation | None: """A model should have an owner defined.""" if "owner" not in model.meta: return RuleViolation("Define the owner of the model in the meta section.") @@ -33,7 +33,7 @@ def has_owner(self, model: Model) -> RuleViolation | None: @rule() -def has_primary_key(self, model: Model) -> RuleViolation | None: +def has_primary_key(self: Rule, model: Model) -> RuleViolation | None: """A model should have a primary key defined, unless it's a view.""" if not model.config.get("materialized") == "picnic_view": has_pk = False @@ -49,7 +49,7 @@ def has_primary_key(self, model: Model) -> RuleViolation | None: @rule() -def primary_key_has_uniqueness_test(self, model: Model) -> RuleViolation | None: +def primary_key_has_uniqueness_test(self: Rule, model: Model) -> RuleViolation | None: """Primary key columns should have a uniqueness test defined.""" columns_with_pk = [] if model.config.get("materialized") == "view": @@ -66,7 +66,7 @@ def primary_key_has_uniqueness_test(self, model: Model) -> RuleViolation | None: @rule() -def columns_have_description(self, model: Model) -> RuleViolation | None: +def columns_have_description(self: Rule, model: Model) -> RuleViolation | None: """All columns of a model should have a description.""" invalid_column_names = [ column.name for column in model.columns if not column.description @@ -81,7 +81,7 @@ def columns_have_description(self, model: Model) -> RuleViolation | None: @rule(description="A model should have at least one test defined.") -def has_test(self, model: Model) -> RuleViolation | None: +def has_test(self: Rule, model: Model) -> RuleViolation | None: """A model should have at least one model-level or column-level test defined. This does not include singular tests, which are tests defined in a separate .sql