yig/leglib/billdb.py

from .common import Bill, CCEColors, CCEAssemblies
from .parsers import BookParser

from typing import Type, Self
from dataclasses import dataclass

class QueryAny:
    """
    Use this class to indicate an Any match for attributes without an Any attribute.
    """
    pass

class SearchNotSatisified(BaseException):
    pass

class QueryAll:
    pass

class QueryField:
    Any = object()
    Colors = CCEColors
    Assemblies = CCEAssemblies

@dataclass
class BillQuery:
    """
    Holds a query for the BillDB.
    """
    color: CCEColors | QueryField = QueryField.Any
    assembly: CCEAssemblies | QueryField = QueryField.Any
    committee: int | QueryField = QueryField.Any
    year: int | QueryField = QueryField.Any
    subcommittee: str | QueryField = QueryField.Any
    sponsors: str | QueryField = QueryField.Any
    school: str | QueryField = QueryField.Any
    bill_text: str | QueryField = QueryField.Any
    title: str | QueryField = QueryField.Any

    def __post_init__(self):
        self.bill_text_concat = self.bill_text # for search compat reasons

class BillDB:
    def __init__(self):
        self.bills: list[Bill] = []
        self.cache: dict[Bill]

    @staticmethod
    def code_enum_match(bill: Bill, query: BillQuery, attr: str) -> None:
        """
        This is probably very slow. Maybe replace this with a better solution?

        This function replaces repetitive code like this:

        elif bill.assembly != CCEAssemblies.Any:
            if bill.assembly != query.color:
                raise SearchNotSatisified()

        with this:

        self.enum_match(bill, query, "color")

        This is the case with exact_match and string_match, too.
        """

        if query.__getattribute__(attr) == QueryField.Any:
            return

        # check the Any case
        if query.__getattribute__(attr) != bill.code.__getattribute__(attr).__class__.Any:
            # make sure we're not matching
            if bill.code.__getattribute__(attr) != query.__getattribute__(attr):
                raise SearchNotSatisified()

        # if we do match, no exception

    @staticmethod
    def string_match(bill: Bill, query: BillQuery, attr: str) -> None:
        """
        See self.code_enum_match for more info.
        """
        if query.__getattribute__(attr) == QueryField.Any:
            return

        if not query.__getattribute__(attr).lower() in bill.__getattribute__(attr).lower():
            raise SearchNotSatisified()

    def add_conference(self: Self, parser: Type[BookParser]) -> None:
        """
        Type[BookParser] -> any subclass of BookParser
        """

        # this works because each BookParser must insert its self.confname into its self.bills[i].code.conference field.
        self.bills += parser.bills

    def search(self: Self, query: BillQuery | QueryAll) -> list[Bill]:
        if query == QueryAll:
            return self.bills
        results = []
        for bill in self.bills:
            try:
                self.code_enum_match(bill, query, "color")
                self.code_enum_match(bill, query, "assembly")

                if not query.committee == QueryField.Any:
                    if not query.committee == bill.code.committee:
                        raise SearchNotSatisified()

                if not query.committee == QueryField.Any:
                    if not query.year == bill.code.year:
                        raise SearchNotSatisified()

                self.string_match(bill, query, "subcommittee")
                self.string_match(bill, query, "sponsors")
                self.string_match(bill, query, "school")
                self.string_match(bill, query, "bill_text_concat")
                self.string_match(bill, query, "title")

            except SearchNotSatisified:
                continue
            results.append(bill)

        return results