grocery-manager/app/data/util.py

#
# Copyright (c) Daniel Sheffield 2023
#
# All rights reserved
#
# THIS SOFTWARE IS PROVIDED AS IS WITHOUT WARRANTY
from typing import List, Set, Tuple, Iterable
from psycopg.sql import (
    Identifier,
    Literal,
    SQL,
    Composable,
)


def get_include_exclude(value) -> Tuple[List[str], List[str]]:
    value = value or ''
    include, exclude = [], []
    if isinstance(value, (list, tuple)):
        for v in value:
            inc, exc = get_include_exclude(v)
            include.extend(inc)
            exclude.extend(exc)
    else:
        inc, exc, *_ = [
            *map(lambda x: x.split('|') if x else [], value.split('!')), []
        ]
        include.extend(inc)
        exclude.extend(exc)
    return list(set(include)), list(set(exclude))


def get_where_include_exclude(
        table: str,
        col: str,
        include: Iterable[str],
        exclude: Iterable[str]
) -> SQL:
    return SQL("""
  ({identifier} = ANY({include}) OR ARRAY[]::text[] @> {include}::text[])
AND
  (NOT {identifier} = ANY({exclude}) OR {identifier} IS NULL)
""").format(
        identifier=Identifier(table, col),
        include=Literal(include),
        exclude=Literal(exclude)
    )

def get_select(alias_to_sql: dict[str,Composable]) -> Composable:
    select = SQL(""",
    """).join([
        SQL(' ').join([
            v, SQL('AS'), Identifier(k)
        ]) for k, v in alias_to_sql.items()
    ])
    return SQL("""
    """).join([SQL("SELECT"), *select])


def get_from(
    base: str,
    table_to_join_on: dict[Tuple[str, Tuple[str,str]]]
) -> Composable:
    joins = [
        SQL("{table} ON {table_column} = {other_column}").format(
            table=Identifier(table),
            table_column=Identifier(table, table_column),
            other_column=Identifier(*other_column) if isinstance(other_column,tuple) else Identifier(other_column)
        ) for table, (table_column, other_column) in table_to_join_on.items()
    ]
    return SQL('').join([SQL("""FROM {base}
LEFT JOIN """).format(base=Identifier(base)),
        SQL("""
LEFT JOIN """).join(joins)])


def get_groupby(alias_to_sql: dict[str, Composable], formatter=None) -> Composable:
    groupby = SQL(""",
""").join([
        formatter(v) if formatter is not None and isinstance(
            v, SQL
        ) else v for k, v in alias_to_sql.items()
    ])
    return SQL("""
""").join([SQL("GROUP BY"), *groupby])