Skip to content

Commit

Permalink
Create rule registry and discovery
Browse files Browse the repository at this point in the history
  • Loading branch information
matthieucan committed Mar 5, 2024
1 parent 7647ddf commit 47282ba
Show file tree
Hide file tree
Showing 8 changed files with 137 additions and 0 deletions.
16 changes: 16 additions & 0 deletions src/dbt_score/definitions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
"""dbt-score definitions."""


from functools import wraps
from typing import Any, Callable


def rule(func: Callable[..., None]) -> Callable[..., None]:
"""Wrapper to create a rule."""

@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
return func(*args, **kwargs)

wrapper._is_rule = True # type: ignore
return wrapper
11 changes: 11 additions & 0 deletions src/dbt_score/exceptions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
"""dbt-score exceptions."""


class DuplicatedRuleException(Exception):
"""Two rules with the same name are defined."""

def __init__(self, rule_name: str):
"""Instantiate exception."""
super().__init__(
f"Rule {rule_name} is defined twice. Rules must have unique names."
)
56 changes: 56 additions & 0 deletions src/dbt_score/registry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
"""Rule registry.
This module implements rule discovery.
"""

import importlib
import pkgutil
from typing import Callable, Iterator

from dbt_score.exceptions import DuplicatedRuleException

THIRD_PARTY_RULES_NAMESPACE = "dbt_score_rules"


class RuleRegistry:
"""A container for configured rules."""

def __init__(self) -> None:
"""Instantiate a rule registry."""
self._rules: dict[str, Callable[[], None]] = {}

@property
def rules(self) -> dict[str, Callable[[], None]]:
"""Get all rules."""
return self._rules

def _walk_packages(self, namespace_name: str) -> Iterator[str]:
"""Walk packages and sub-packages recursively."""
try:
namespace = importlib.import_module(namespace_name)
except ImportError: # no custom rule in Python path
return

def onerror(module_name: str) -> None:
print(f"Failed to import {module_name}.")

for package in pkgutil.walk_packages(namespace.__path__, onerror=onerror):
yield f"{namespace_name}.{package.name}"
if package.ispkg:
yield from self._walk_packages(f"{namespace_name}.{package.name}")

def _load(self, namespace_name: str) -> None:
"""Load rules found in a given namespace."""
for module_name in self._walk_packages(namespace_name):
module = importlib.import_module(module_name)
for obj_name in dir(module):
obj = module.__dict__[obj_name]
if getattr(obj, "_is_rule", False):
if obj_name in self.rules:
raise DuplicatedRuleException(obj_name)
self._rules[obj_name] = obj

def load_all(self) -> None:
"""Load all rules, core and third-party."""
self._load("dbt_score.rules")
self._load(THIRD_PARTY_RULES_NAMESPACE)
8 changes: 8 additions & 0 deletions src/dbt_score/rules/example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
"""Placeholder for example rules."""

from dbt_score.definitions import rule


@rule
def rule_example() -> None:
"""An example rule."""
9 changes: 9 additions & 0 deletions tests/rules/example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
"""Example rules."""


from dbt_score.definitions import rule


@rule
def rule_test_example():
"""An example rule."""
1 change: 1 addition & 0 deletions tests/rules/nested/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""Nested package for testing rule discovery."""
9 changes: 9 additions & 0 deletions tests/rules/nested/example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
"""Example rules."""


from dbt_score.definitions import rule


@rule
def rule_test_nested_example():
"""An example rule."""
27 changes: 27 additions & 0 deletions tests/test_rule_registry.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
"""Unit tests for the rule registry."""

import pytest
from dbt_score.exceptions import DuplicatedRuleException
from dbt_score.registry import RuleRegistry


def test_rule_registry_discovery():
"""Ensure rules can be found in a given namespace recursively."""
r = RuleRegistry()
r._load("tests.rules")
assert sorted(r.rules.keys()) == ["rule_test_example", "rule_test_nested_example"]


def test_rule_registry_no_duplicates():
"""Ensure no duplicate rule names can coexist."""
r = RuleRegistry()
r._load("tests.rules")
with pytest.raises(DuplicatedRuleException):
r._load("tests.rules")


def test_rule_registry_core_rules():
"""Ensure core rules are automatically discovered."""
r = RuleRegistry()
r.load_all()
assert len(r.rules) > 0

0 comments on commit 47282ba

Please sign in to comment.