codekingpro/portable-devtools

1"""2psycopg row factories3"""4 5# Copyright (C) 2021 The Psycopg Team6 7import functools8from typing import Any, Callable, Dict, List, Optional, NamedTuple, NoReturn9from typing import TYPE_CHECKING, Sequence, Tuple, Type, TypeVar10from collections import namedtuple11from typing_extensions import TypeAlias12 13from . import pq14from . import errors as e15from ._compat import Protocol16from ._encodings import _as_python_identifier17 18if TYPE_CHECKING:19    from .cursor import BaseCursor, Cursor20    from .cursor_async import AsyncCursor21    from psycopg.pq.abc import PGresult22 23COMMAND_OK = pq.ExecStatus.COMMAND_OK24TUPLES_OK = pq.ExecStatus.TUPLES_OK25SINGLE_TUPLE = pq.ExecStatus.SINGLE_TUPLE26 27T = TypeVar("T", covariant=True)28 29# Row factories30 31Row = TypeVar("Row", covariant=True)32 33 34class RowMaker(Protocol[Row]):35    """36    Callable protocol taking a sequence of value and returning an object.37 38    The sequence of value is what is returned from a database query, already39    adapted to the right Python types. The return value is the object that your40    program would like to receive: by default (`tuple_row()`) it is a simple41    tuple, but it may be any type of object.42 43    Typically, `!RowMaker` functions are returned by `RowFactory`.44    """45 46    def __call__(self, __values: Sequence[Any]) -> Row:47        ...48 49 50class RowFactory(Protocol[Row]):51    """52    Callable protocol taking a `~psycopg.Cursor` and returning a `RowMaker`.53 54    A `!RowFactory` is typically called when a `!Cursor` receives a result.55    This way it can inspect the cursor state (for instance the56    `~psycopg.Cursor.description` attribute) and help a `!RowMaker` to create57    a complete object.58 59    For instance the `dict_row()` `!RowFactory` uses the names of the column to60    define the dictionary key and returns a `!RowMaker` function which would61    use the values to create a dictionary for each record.62    """63 64    def __call__(self, __cursor: "Cursor[Any]") -> RowMaker[Row]:65        ...66 67 68class AsyncRowFactory(Protocol[Row]):69    """70    Like `RowFactory`, taking an async cursor as argument.71    """72 73    def __call__(self, __cursor: "AsyncCursor[Any]") -> RowMaker[Row]:74        ...75 76 77class BaseRowFactory(Protocol[Row]):78    """79    Like `RowFactory`, taking either type of cursor as argument.80    """81 82    def __call__(self, __cursor: "BaseCursor[Any, Any]") -> RowMaker[Row]:83        ...84 85 86TupleRow: TypeAlias = Tuple[Any, ...]87"""88An alias for the type returned by `tuple_row()` (i.e. a tuple of any content).89"""90 91 92DictRow: TypeAlias = Dict[str, Any]93"""94An alias for the type returned by `dict_row()`95 96A `!DictRow` is a dictionary with keys as string and any value returned by the97database.98"""99 100 101def tuple_row(cursor: "BaseCursor[Any, Any]") -> "RowMaker[TupleRow]":102    r"""Row factory to represent rows as simple tuples.103 104    This is the default factory, used when `~psycopg.Connection.connect()` or105    `~psycopg.Connection.cursor()` are called without a `!row_factory`106    parameter.107 108    """109    # Implementation detail: make sure this is the tuple type itself, not an110    # equivalent function, because the C code fast-paths on it.111    return tuple112 113 114def dict_row(cursor: "BaseCursor[Any, Any]") -> "RowMaker[DictRow]":115    """Row factory to represent rows as dictionaries.116 117    The dictionary keys are taken from the column names of the returned columns.118    """119    names = _get_names(cursor)120    if names is None:121        return no_result122 123    def dict_row_(values: Sequence[Any]) -> Dict[str, Any]:124        return dict(zip(names, values))125 126    return dict_row_127 128 129def namedtuple_row(130    cursor: "BaseCursor[Any, Any]",131) -> "RowMaker[NamedTuple]":132    """Row factory to represent rows as `~collections.namedtuple`.133 134    The field names are taken from the column names of the returned columns,135    with some mangling to deal with invalid names.136    """137    res = cursor.pgresult138    if not res:139        return no_result140 141    nfields = _get_nfields(res)142    if nfields is None:143        return no_result144 145    nt = _make_nt(cursor._encoding, *(res.fname(i) for i in range(nfields)))146    return nt._make147 148 149@functools.lru_cache(512)150def _make_nt(enc: str, *names: bytes) -> Type[NamedTuple]:151    snames = tuple(_as_python_identifier(n.decode(enc)) for n in names)152    return namedtuple("Row", snames)  # type: ignore[return-value]153 154 155def class_row(cls: Type[T]) -> BaseRowFactory[T]:156    r"""Generate a row factory to represent rows as instances of the class `!cls`.157 158    The class must support every output column name as a keyword parameter.159 160    :param cls: The class to return for each row. It must support the fields161        returned by the query as keyword arguments.162    :rtype: `!Callable[[Cursor],` `RowMaker`\[~T]]163    """164 165    def class_row_(cursor: "BaseCursor[Any, Any]") -> "RowMaker[T]":166        names = _get_names(cursor)167        if names is None:168            return no_result169 170        def class_row__(values: Sequence[Any]) -> T:171            return cls(**dict(zip(names, values)))172 173        return class_row__174 175    return class_row_176 177 178def args_row(func: Callable[..., T]) -> BaseRowFactory[T]:179    """Generate a row factory calling `!func` with positional parameters for every row.180 181    :param func: The function to call for each row. It must support the fields182        returned by the query as positional arguments.183    """184 185    def args_row_(cur: "BaseCursor[Any, T]") -> "RowMaker[T]":186        def args_row__(values: Sequence[Any]) -> T:187            return func(*values)188 189        return args_row__190 191    return args_row_192 193 194def kwargs_row(func: Callable[..., T]) -> BaseRowFactory[T]:195    """Generate a row factory calling `!func` with keyword parameters for every row.196 197    :param func: The function to call for each row. It must support the fields198        returned by the query as keyword arguments.199    """200 201    def kwargs_row_(cursor: "BaseCursor[Any, T]") -> "RowMaker[T]":202        names = _get_names(cursor)203        if names is None:204            return no_result205 206        def kwargs_row__(values: Sequence[Any]) -> T:207            return func(**dict(zip(names, values)))208 209        return kwargs_row__210 211    return kwargs_row_212 213 214def no_result(values: Sequence[Any]) -> NoReturn:215    """A `RowMaker` that always fail.216 217    It can be used as return value for a `RowFactory` called with no result.218    Note that the `!RowFactory` *will* be called with no result, but the219    resulting `!RowMaker` never should.220    """221    raise e.InterfaceError("the cursor doesn't have a result")222 223 224def _get_names(cursor: "BaseCursor[Any, Any]") -> Optional[List[str]]:225    res = cursor.pgresult226    if not res:227        return None228 229    nfields = _get_nfields(res)230    if nfields is None:231        return None232 233    enc = cursor._encoding234    return [235        res.fname(i).decode(enc) for i in range(nfields)  # type: ignore[union-attr]236    ]237 238 239def _get_nfields(res: "PGresult") -> Optional[int]:240    """241    Return the number of columns in a result, if it returns tuples else None242 243    Take into account the special case of results with zero columns.244    """245    nfields = res.nfields246 247    if (248        res.status == TUPLES_OK249        or res.status == SINGLE_TUPLE250        # "describe" in named cursors251        or (res.status == COMMAND_OK and nfields)252    ):253        return nfields254    else:255        return None256