SQLAlchemy 2.1.0rc1__cp315-cp315-win_amd64.whl

sqlalchemy/engine/_result_cy.py

@@ -0,0 +1,711 @@
+# engine/_result_cy.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: disable-error-code="misc,no-redef,type-arg,untyped-decorator"
+from __future__ import annotations
+
+from collections.abc import Callable
+from collections.abc import Iterator
+from collections.abc import Sequence
+from enum import Enum
+import operator
+from typing import Any
+from typing import cast
+from typing import Generic
+from typing import Literal
+from typing import overload
+from typing import TYPE_CHECKING
+from typing import TypeVar
+from typing import Union
+
+from .row import Row
+from .row import RowMapping
+from .. import exc
+from ..util import HasMemoized_ro_memoized_attribute
+from ..util.typing import Self
+from ..util.typing import TupleAny
+from ..util.typing import Unpack
+
+if TYPE_CHECKING:
+    from .result import _ProcessorsType
+    from .result import Result
+    from .result import ResultMetaData
+
+# START GENERATED CYTHON IMPORT
+# This section is automatically generated by the script tools/cython_imports.py
+try:
+    # NOTE: the cython compiler needs this "import cython" in the file, it
+    # can't be only "from sqlalchemy.util import cython" with the fallback
+    # in that module
+    import cython
+except ModuleNotFoundError:
+    from sqlalchemy.util import cython
+
+
+def _is_compiled() -> bool:
+    """Utility function to indicate if this module is compiled or not."""
+    return cython.compiled  # type: ignore[no-any-return,unused-ignore]
+
+
+# END GENERATED CYTHON IMPORT
+
+if cython.compiled:
+    from cython.cimports.cpython import Py_INCREF
+    from cython.cimports.cpython import PyList_New
+    from cython.cimports.cpython import PyList_SET_ITEM
+    from cython.cimports.cpython import PyTuple_New
+    from cython.cimports.cpython import PyTuple_SET_ITEM
+
+
+_RowData = Row[Unpack[TupleAny]] | RowMapping | Any
+"""A generic form of "row" that accommodates for the different kinds of
+"rows" that different result objects return, including row, row mapping, and
+scalar values"""
+_R = TypeVar("_R", bound=_RowData)
+_T = TypeVar("_T", bound=Any)
+
+_InterimRowType = Union[_R, TupleAny]
+"""a catchall "anything" kind of return type that can be applied
+across all the result types
+
+"""
+
+_UniqueFilterType = Callable[[Any], Any]
+_UniqueFilterStateType = tuple[set[Any], _UniqueFilterType | None]
+
+_FLAG_SIMPLE = cython.declare(cython.const[cython.char], 0)
+_FLAG_SCALAR_TO_TUPLE = cython.declare(cython.const[cython.char], 1)
+_FLAG_TUPLE_FILTER = cython.declare(cython.const[cython.char], 2)
+
+
+# a symbol that indicates to internal Result methods that
+# "no row is returned".  We can't use None for those cases where a scalar
+# filter is applied to rows.
+class _NoRow(Enum):
+    _NO_ROW = 0
+
+
+_NO_ROW = _NoRow._NO_ROW
+
+
+class BaseResultInternal(Generic[_R]):
+    __slots__ = ()
+
+    _real_result: Result[Unpack[TupleAny]] | None = None
+    _generate_rows: bool = True
+    _row_logging_fn: Callable[[Any], Any] | None
+
+    _unique_filter_state: _UniqueFilterStateType | None = None
+    _post_creational_filter: Callable[[Any], Any] | None = None
+
+    _metadata: ResultMetaData
+
+    _source_supports_scalars: bool
+    _yield_per: int | None
+
+    def _fetchiter_impl(
+        self,
+    ) -> Iterator[_InterimRowType[Row[Unpack[TupleAny]]]]:
+        raise NotImplementedError()
+
+    def _fetchone_impl(
+        self, hard_close: bool = False
+    ) -> _InterimRowType[Row[Unpack[TupleAny]]] | None:
+        raise NotImplementedError()
+
+    def _fetchmany_impl(
+        self, size: int | None = None
+    ) -> list[_InterimRowType[Row[Unpack[TupleAny]]]]:
+        raise NotImplementedError()
+
+    def _fetchall_impl(
+        self,
+    ) -> list[_InterimRowType[Row[Unpack[TupleAny]]]]:
+        raise NotImplementedError()
+
+    def _soft_close(self, hard: bool = False) -> None:
+        raise NotImplementedError()
+
+    @HasMemoized_ro_memoized_attribute
+    def _row_getter(
+        self,
+    ) -> tuple[
+        Callable[..., _R] | None,
+        Callable[..., Sequence[_R]] | None,
+        Callable[..., Sequence[Any]] | None,
+    ]:
+        """Return a three-tuple of callables
+        ``(single_row, many_rows, interim_rows)``.
+
+        ``single_row`` and ``many_rows`` produce :class:`.Row` objects for
+        one input row or a sequence of them.  ``interim_rows`` produces
+        rows in "interim" form for a sequence of input rows: plain tuples
+        with result processors and tuple filters applied but no
+        :class:`.Row` construction, falling back to ``many_rows`` when
+        Row objects are required (row logging, scalar sources).
+
+        """
+        real_result = self if self._real_result is None else self._real_result
+
+        metadata = self._metadata
+        tuple_filters = metadata._tuplefilter
+        flag: cython.char = _FLAG_SIMPLE
+
+        if real_result._source_supports_scalars:
+            if not self._generate_rows:
+                return None, None, None
+            else:
+                flag = _FLAG_SCALAR_TO_TUPLE
+        elif tuple_filters is not None:
+            flag = _FLAG_TUPLE_FILTER
+
+        processors: tuple
+        proc_valid: tuple
+
+        if metadata._effective_processors is not None:
+            ep = metadata._effective_processors
+            if flag == _FLAG_TUPLE_FILTER:
+                ep = tuple_filters(ep)
+
+            processors = tuple(ep)
+            proc_valid = tuple(
+                [i for i, p in enumerate(processors) if p is not None]
+            )
+        else:
+            processors = ()
+            proc_valid = ()
+
+        proc_size: cython.Py_ssize_t = len(processors)
+        log_row = real_result._row_logging_fn
+        has_log_row: cython.bint = log_row is not None
+
+        key_to_index = metadata._key_to_index
+        _Row = Row
+
+        if flag == _FLAG_SIMPLE and proc_size == 0 and not has_log_row:
+            # just build the rows
+
+            def single_row_simple(input_row: Sequence[Any], /) -> Row:
+                return _Row(metadata, None, key_to_index, input_row)
+
+            if cython.compiled:
+
+                def many_rows_simple(rows: Sequence[Any], /) -> list[Any]:
+                    size: cython.Py_hash_t = len(rows)
+                    i: cython.Py_ssize_t
+                    result: list = PyList_New(size)
+                    for i in range(size):
+                        row: object = _Row(
+                            metadata, None, key_to_index, rows[i]
+                        )
+                        Py_INCREF(row)
+                        PyList_SET_ITEM(result, i, row)
+                    return result
+
+            else:
+
+                def many_rows_simple(rows: Sequence[Any], /) -> list[Any]:
+                    return [
+                        _Row(metadata, None, key_to_index, row) for row in rows
+                    ]
+
+            def interim_rows_simple(rows: Sequence[Any], /) -> Sequence[Any]:
+                # interim rows must be tuples so that position-based
+                # getters and slices apply; most DBAPIs return tuples
+                # already
+                if rows and type(rows[0]) is not tuple:
+                    return [tuple(row) for row in rows]
+                return rows
+
+            return single_row_simple, many_rows_simple, interim_rows_simple  # type: ignore[return-value] # noqa: E501
+
+        first_row: cython.bint = True
+
+        def single_row(input_row: Sequence[Any], /) -> Row:
+            nonlocal first_row
+
+            if flag == _FLAG_SCALAR_TO_TUPLE:
+                input_row = (input_row,)
+            elif flag == _FLAG_TUPLE_FILTER:
+                input_row = tuple_filters(input_row)
+
+            if proc_size != 0:
+                if first_row:
+                    first_row = False
+                    assert len(input_row) == proc_size
+                input_row = _apply_processors(
+                    processors, proc_size, proc_valid, input_row
+                )
+
+            row: Row = _Row(metadata, None, key_to_index, input_row)
+            if has_log_row:
+                row = log_row(row)
+            return row
+
+        if cython.compiled:
+
+            def many_rows(rows: Sequence[Any], /) -> list[Any]:
+                size: cython.Py_hash_t = len(rows)
+                i: cython.Py_ssize_t
+                result: list = PyList_New(size)
+                for i in range(size):
+                    row: object = single_row(rows[i])
+                    Py_INCREF(row)
+                    PyList_SET_ITEM(result, i, row)
+                return result
+
+        else:
+
+            def many_rows(rows: Sequence[Any], /) -> list[Any]:
+                return [single_row(row) for row in rows]
+
+        if flag == _FLAG_SCALAR_TO_TUPLE or has_log_row:
+            # scalar-source rows have no tuple form; row logging requires
+            # the Row objects themselves
+            interim_rows = many_rows
+        else:
+
+            def single_interim_row(input_row: Sequence[Any], /) -> Any:
+                if flag == _FLAG_TUPLE_FILTER:
+                    input_row = tuple_filters(input_row)
+                if proc_size != 0:
+                    input_row = _apply_processors(
+                        processors, proc_size, proc_valid, input_row
+                    )
+                elif type(input_row) is not tuple:
+                    input_row = tuple(input_row)
+                return input_row
+
+            if cython.compiled:
+
+                def interim_rows(rows: Sequence[Any], /) -> Sequence[Any]:
+                    size: cython.Py_hash_t = len(rows)
+                    i: cython.Py_ssize_t
+                    result: list = PyList_New(size)
+                    for i in range(size):
+                        row: object = single_interim_row(rows[i])
+                        Py_INCREF(row)
+                        PyList_SET_ITEM(result, i, row)
+                    return result
+
+            else:
+
+                def interim_rows(rows: Sequence[Any], /) -> Sequence[Any]:
+                    return [single_interim_row(row) for row in rows]
+
+        return single_row, many_rows, interim_rows  # type: ignore[return-value] # noqa: E501
+
+    @HasMemoized_ro_memoized_attribute
+    def _iterator_getter(self) -> Callable[[], Iterator[_R]]:
+        make_row = self._row_getter[0]
+
+        post_creational_filter = self._post_creational_filter
+
+        if self._unique_filter_state is not None:
+            uniques: set
+            uniques, strategy = self._unique_strategy
+
+            def iterrows() -> Iterator[_R]:
+                for raw_row in self._fetchiter_impl():
+                    row = (
+                        make_row(raw_row) if make_row is not None else raw_row
+                    )
+                    hashed = strategy(row) if strategy is not None else row
+                    if hashed in uniques:
+                        continue
+                    uniques.add(hashed)
+                    if post_creational_filter is not None:
+                        row = post_creational_filter(row)
+                    yield row
+
+        else:
+
+            def iterrows() -> Iterator[_R]:
+                for raw_row in self._fetchiter_impl():
+                    row = (
+                        make_row(raw_row) if make_row is not None else raw_row
+                    )
+                    if post_creational_filter is not None:
+                        row = post_creational_filter(row)
+                    yield row
+
+        return iterrows
+
+    def _raw_all_tuples(self) -> Sequence[tuple[Any, ...]]:
+        """Return all remaining rows as plain processed tuples, as produced
+        by the ``interim_rows`` element of :attr:`._row_getter`;
+        :class:`.Row` objects are produced instead where required (row
+        logging, scalar sources).
+
+        Used by ORM loading, whose row getters are position-based and
+        accept any tuple-like row.
+
+        Uniquing and post-creational filters are not applied here, so
+        may not be present on the :class:`.Result` when this method is
+        used.
+
+        """
+        assert self._unique_filter_state is None
+        assert self._post_creational_filter is None
+        interim_rows = self._row_getter[2]
+        assert interim_rows is not None
+        rows = self._fetchall_impl()
+        return interim_rows(rows)
+
+    def _allrows(self) -> Sequence[_R]:
+        post_creational_filter = self._post_creational_filter
+
+        make_rows = self._row_getter[1]
+
+        rows = self._fetchall_impl()
+        made_rows: Sequence[_InterimRowType[_R]]
+        if make_rows is not None:
+            made_rows = make_rows(rows)
+        else:
+            made_rows = rows
+
+        interim_rows: Sequence[_R]
+
+        if self._unique_filter_state is not None:
+            uniques: set
+            uniques, strategy = self._unique_strategy
+            interim_rows = _apply_unique_strategy(
+                made_rows, [], uniques, strategy
+            )
+        else:
+            interim_rows = made_rows  # type: ignore[assignment]
+
+        if post_creational_filter is not None:
+            interim_rows = [
+                post_creational_filter(row) for row in interim_rows
+            ]
+        return interim_rows
+
+    @HasMemoized_ro_memoized_attribute
+    def _onerow_getter(
+        self,
+    ) -> Callable[[Self], Literal[_NoRow._NO_ROW] | _R]:
+        make_row = self._row_getter[0]
+
+        post_creational_filter = self._post_creational_filter
+
+        if self._unique_filter_state is not None:
+            uniques: set
+            uniques, strategy = self._unique_strategy
+
+            def onerow(self: Self) -> Literal[_NoRow._NO_ROW] | _R:
+                while True:
+                    row = self._fetchone_impl()
+                    if row is None:
+                        return _NO_ROW
+                    else:
+                        obj: _InterimRowType[Any] = (
+                            make_row(row) if make_row is not None else row
+                        )
+                        hashed = strategy(obj) if strategy is not None else obj
+                        if hashed in uniques:
+                            continue
+                        uniques.add(hashed)
+                        if post_creational_filter is not None:
+                            obj = post_creational_filter(obj)
+                        return obj  # type: ignore[return-value]
+
+        else:
+
+            def onerow(self: Self) -> Literal[_NoRow._NO_ROW] | _R:
+                row = self._fetchone_impl()
+                if row is None:
+                    return _NO_ROW
+                else:
+                    interim_row: _InterimRowType[Any] = (
+                        make_row(row) if make_row is not None else row
+                    )
+                    if post_creational_filter is not None:
+                        interim_row = post_creational_filter(interim_row)
+                    return interim_row  # type: ignore[return-value]
+
+        return onerow
+
+    @HasMemoized_ro_memoized_attribute
+    def _manyrow_getter(self) -> Callable[[Self, int | None], Sequence[_R]]:
+        make_rows = self._row_getter[1]
+        real_result = self if self._real_result is None else self._real_result
+        yield_per = real_result._yield_per
+
+        post_creational_filter = self._post_creational_filter
+
+        if self._unique_filter_state:
+            uniques: set
+            uniques, strategy = self._unique_strategy
+
+            def manyrows(self: Self, num: int | None, /) -> Sequence[_R]:
+                made_rows: Sequence[Any]
+                collect: list[_R] = []
+
+                _manyrows = self._fetchmany_impl
+
+                if num is None:
+                    # if None is passed, we don't know the default
+                    # manyrows number, DBAPI has this as cursor.arraysize
+                    # different DBAPIs / fetch strategies may be different.
+                    # do a fetch to find what the number is.  if there are
+                    # only fewer rows left, then it doesn't matter.
+                    if yield_per:
+                        num_required = num = yield_per
+                    else:
+                        rows = _manyrows()
+                        num = len(rows)
+                        made_rows = (
+                            rows if make_rows is None else make_rows(rows)
+                        )
+                        _apply_unique_strategy(
+                            made_rows, collect, uniques, strategy
+                        )
+                        num_required = num - len(collect)
+                else:
+                    num_required = num
+
+                assert num is not None
+
+                while num_required:
+                    rows = _manyrows(num_required)
+                    if not rows:
+                        break
+
+                    made_rows = rows if make_rows is None else make_rows(rows)
+                    _apply_unique_strategy(
+                        made_rows, collect, uniques, strategy
+                    )
+                    num_required = num - len(collect)
+
+                if post_creational_filter is not None:
+                    collect = [post_creational_filter(row) for row in collect]
+                return collect
+
+        else:
+
+            def manyrows(self: Self, num: int | None, /) -> Sequence[_R]:
+                if num is None:
+                    num = yield_per
+
+                rows: Sequence = self._fetchmany_impl(num)
+                if make_rows is not None:
+                    rows = make_rows(rows)
+                if post_creational_filter is not None:
+                    rows = [post_creational_filter(row) for row in rows]
+                return rows
+
+        return manyrows
+
+    @overload
+    def _only_one_row(
+        self: BaseResultInternal[Row[_T, Unpack[TupleAny]]],
+        raise_for_second_row: bool,
+        raise_for_none: bool,
+        scalar: Literal[True],
+    ) -> _T: ...
+
+    @overload
+    def _only_one_row(
+        self,
+        raise_for_second_row: bool,
+        raise_for_none: Literal[True],
+        scalar: bool,
+    ) -> _R: ...
+
+    @overload
+    def _only_one_row(
+        self,
+        raise_for_second_row: bool,
+        raise_for_none: bool,
+        scalar: bool,
+    ) -> _R | None: ...
+
+    def _only_one_row(
+        self,
+        raise_for_second_row: bool,
+        raise_for_none: bool,
+        scalar: bool,
+    ) -> _R | None:
+        onerow = self._fetchone_impl
+
+        row = onerow(hard_close=True)
+        if row is None:
+            if raise_for_none:
+                raise exc.NoResultFound(
+                    "No row was found when one was required"
+                )
+            else:
+                return None
+
+        if scalar and self._source_supports_scalars:
+            self._generate_rows = False
+            make_row = None
+        else:
+            make_row = self._row_getter[0]
+
+        try:
+            row = make_row(row) if make_row is not None else row  # type: ignore[assignment] # noqa: E501
+        except:
+            self._soft_close(hard=True)
+            raise
+
+        if raise_for_second_row:
+            if self._unique_filter_state:
+                # for no second row but uniqueness, need to essentially
+                # consume the entire result :(
+                strategy = self._unique_strategy[1]
+
+                existing_row_hash = (
+                    strategy(row) if strategy is not None else row
+                )
+
+                while True:
+                    next_row: Any = onerow(hard_close=True)
+                    if next_row is None:
+                        next_row = _NO_ROW
+                        break
+
+                    try:
+                        next_row = (
+                            make_row(next_row)
+                            if make_row is not None
+                            else next_row
+                        )
+
+                        if strategy is not None:
+                            # assert next_row is not _NO_ROW
+                            if existing_row_hash == strategy(next_row):
+                                continue
+                        elif row == next_row:
+                            continue
+                        # here, we have a row and it's different
+                        break
+                    except:
+                        self._soft_close(hard=True)
+                        raise
+            else:
+                next_row = onerow(hard_close=True)
+                if next_row is None:
+                    next_row = _NO_ROW
+
+            if next_row is not _NO_ROW:
+                self._soft_close(hard=True)
+                raise exc.MultipleResultsFound(
+                    "Multiple rows were found when exactly one was required"
+                    if raise_for_none
+                    else "Multiple rows were found when one or none "
+                    "was required"
+                )
+        else:
+            next_row = _NO_ROW
+            # if we checked for second row then that would have
+            # closed us :)
+            self._soft_close(hard=True)
+
+        if not scalar:
+            post_creational_filter = self._post_creational_filter
+            if post_creational_filter is not None:
+                row = post_creational_filter(row)
+
+        if scalar and make_row is not None:
+            return row[0]  # type: ignore[no-any-return]
+        else:
+            return row  # type: ignore[return-value]
+
+    def _iter_impl(self) -> Iterator[_R]:
+        return self._iterator_getter()
+
+    def _next_impl(self) -> _R:
+        row = self._onerow_getter(self)
+        if row is _NO_ROW:
+            raise StopIteration()
+        else:
+            return row
+
+    @HasMemoized_ro_memoized_attribute
+    def _unique_strategy(self) -> _UniqueFilterStateType:
+        assert self._unique_filter_state is not None
+        uniques, strategy = self._unique_filter_state
+
+        if (
+            strategy is None
+            and self._metadata._create_unique_filters is not None
+        ):
+            real_result = cast(
+                "Result[Any]",
+                self if self._real_result is None else self._real_result,
+            )
+            filters = self._metadata._create_unique_filters(real_result)
+            if (
+                real_result._source_supports_scalars
+                and not self._generate_rows
+            ):
+                strategy = filters[0]
+            else:
+                if self._metadata._tuplefilter is not None:
+                    filters = self._metadata._tuplefilter(filters)
+
+                strategy = operator.methodcaller("_filter_on_values", filters)
+        return uniques, strategy
+
+
+if cython.compiled:
+
+    @cython.inline
+    @cython.cfunc
+    @cython.wraparound(False)
+    @cython.boundscheck(False)
+    def _apply_processors(
+        proc: tuple,
+        proc_size: cython.Py_ssize_t,
+        proc_valid: object,  # used only by python impl
+        data: Sequence,
+    ) -> tuple[Any, ...]:
+        res: tuple = PyTuple_New(proc_size)
+        i: cython.Py_ssize_t
+        for i in range(proc_size):
+            p = proc[i]
+            if p is not None:
+                value = p(data[i])
+            else:
+                value = data[i]
+            Py_INCREF(value)
+            PyTuple_SET_ITEM(res, i, value)
+        return res
+
+else:
+
+    def _apply_processors(
+        proc: _ProcessorsType,
+        proc_size: int,  # used only by cython impl
+        proc_valid: tuple[int, ...],
+        data: Sequence[Any],
+    ) -> tuple[Any, ...]:
+        res = list(data)
+        for i in proc_valid:
+            res[i] = proc[i](res[i])
+        return tuple(res)
+
+
+@cython.inline
+@cython.cfunc
+def _apply_unique_strategy(
+    rows: Sequence[Any],
+    destination: list[Any],
+    uniques: set[Any],
+    strategy: Callable[[Any], Any] | None,
+) -> list[Any]:
+    i: cython.Py_ssize_t
+    has_strategy: cython.bint = strategy is not None
+    for i in range(len(rows)):
+        row = rows[i]
+        hashed = strategy(row) if has_strategy else row
+        if hashed in uniques:
+            continue
+        uniques.add(hashed)
+        destination.append(row)
+    return destination