From 0b236dd9ace2834ab680e3726b4cfd2d4b8a34b1 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 12:11:59 +0900 Subject: [PATCH 1/6] Apply dict_type only to the dict cursor it was given to DictCursor, AsyncDictCursor and AioDictCursor assigned dict_type to the class attribute of the shared dict result set class, so one cursor's dict_type changed the row type of every dict cursor in the process. When dict_type is given, the cursor now uses its own subclass of the result set class that carries it, and the shared class is left unchanged. Closes #923 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 13 ++++++++++--- pyathena/async_cursor.py | 13 ++++++++++--- pyathena/cursor.py | 13 ++++++++++--- tests/pyathena/aio/test_cursor.py | 9 +++++++++ tests/pyathena/test_async_cursor.py | 9 +++++++++ tests/pyathena/test_cursor.py | 9 +++++++++ 6 files changed, 57 insertions(+), 9 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 89e7c5aa6..e34e9871c 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -251,10 +251,17 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``AioCursor.__init__``. If they include - ``dict_type``, it is also assigned to the class attribute - ``AthenaAioDictResultSet.dict_type``, the type used to build each row. + ``dict_type``, it is the type used to build each row of this + cursor's result sets; other cursors are not affected. """ super().__init__(**kwargs) self._result_set_class = AthenaAioDictResultSet if "dict_type" in kwargs: - AthenaAioDictResultSet.dict_type = kwargs["dict_type"] + self._result_set_class = cast( + type[AthenaAioDictResultSet], + type( + AthenaAioDictResultSet.__name__, + (AthenaAioDictResultSet,), + {"dict_type": kwargs["dict_type"]}, + ), + ) diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 3363901c7..a0766b484 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -340,10 +340,17 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``AsyncCursor.__init__``. If they include - ``dict_type``, it is also assigned to the class attribute - ``AthenaDictResultSet.dict_type``, the type used to build each row. + ``dict_type``, it is the type used to build each row of this + cursor's result sets; other cursors are not affected. """ super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet if "dict_type" in kwargs: - AthenaDictResultSet.dict_type = kwargs["dict_type"] + self._result_set_class = cast( + type[AthenaDictResultSet], + type( + AthenaDictResultSet.__name__, + (AthenaDictResultSet,), + {"dict_type": kwargs["dict_type"]}, + ), + ) diff --git a/pyathena/cursor.py b/pyathena/cursor.py index d0b6de4f3..f71302858 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -221,10 +221,17 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``Cursor.__init__``. If they include - ``dict_type``, it is also assigned to the class attribute - ``AthenaDictResultSet.dict_type``, the type used to build each row. + ``dict_type``, it is the type used to build each row of this + cursor's result sets; other cursors are not affected. """ super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet if "dict_type" in kwargs: - AthenaDictResultSet.dict_type = kwargs["dict_type"] + self._result_set_class = cast( + type[AthenaDictResultSet], + type( + AthenaDictResultSet.__name__, + (AthenaDictResultSet,), + {"dict_type": kwargs["dict_type"]}, + ), + ) diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 568449568..8990a1a99 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,6 +1,7 @@ import asyncio import re import threading +from collections import OrderedDict from datetime import UTC, datetime, timedelta from unittest.mock import AsyncMock, MagicMock, call, patch @@ -896,3 +897,11 @@ async def test_fetchall(self, aio_dict_cursor): assert await aio_dict_cursor.fetchall() == [{"number_of_rows": 1}] await aio_dict_cursor.execute("SELECT a FROM many_rows ORDER BY a") assert await aio_dict_cursor.fetchall() == [{"a": i} for i in range(10000)] + + async def test_dict_type(self, aio_dict_cursor): + async with aio_dict_cursor.connection.cursor(dict_type=OrderedDict) as ordered_cursor: + await ordered_cursor.execute("SELECT * FROM one_row") + assert type(await ordered_cursor.fetchone()) is OrderedDict + # dict_type of another cursor does not change the row type of this one. + await aio_dict_cursor.execute("SELECT * FROM one_row") + assert type(await aio_dict_cursor.fetchone()) is dict diff --git a/tests/pyathena/test_async_cursor.py b/tests/pyathena/test_async_cursor.py index 1cf4d0d91..710292686 100644 --- a/tests/pyathena/test_async_cursor.py +++ b/tests/pyathena/test_async_cursor.py @@ -1,5 +1,6 @@ import contextlib import time +from collections import OrderedDict from datetime import datetime from random import randint @@ -246,3 +247,11 @@ def test_fetchall(self, async_dict_cursor): query_id, future = async_dict_cursor.execute("SELECT a FROM many_rows ORDER BY a") result_set = future.result() assert result_set.fetchall() == [{"a": i} for i in range(10000)] + + def test_dict_type(self, async_dict_cursor): + with async_dict_cursor.connection.cursor(dict_type=OrderedDict) as ordered_cursor: + _, future = ordered_cursor.execute("SELECT * FROM one_row") + assert type(future.result().fetchone()) is OrderedDict + # dict_type of another cursor does not change the row type of this one. + _, future = async_dict_cursor.execute("SELECT * FROM one_row") + assert type(future.result().fetchone()) is dict diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 720a5b230..e3f440e67 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -7,6 +7,7 @@ import threading import time import uuid +from collections import OrderedDict from concurrent import futures from concurrent.futures.thread import ThreadPoolExecutor from datetime import UTC, date, datetime, timedelta @@ -1876,6 +1877,14 @@ def test_fetchall(self, dict_cursor): dict_cursor.execute("SELECT a FROM many_rows ORDER BY a") assert dict_cursor.fetchall() == [{"a": i} for i in range(10000)] + def test_dict_type(self, dict_cursor): + with dict_cursor.connection.cursor(dict_type=OrderedDict) as ordered_cursor: + ordered_cursor.execute("SELECT * FROM one_row") + assert type(ordered_cursor.fetchone()) is OrderedDict + # dict_type of another cursor does not change the row type of this one. + dict_cursor.execute("SELECT * FROM one_row") + assert type(dict_cursor.fetchone()) is dict + def test_null_vs_empty_string(self, dict_cursor): """ DictCursor should properly distinguish NULL from empty string. From 27626643b1770bd9cc143c3a5f1486c018475f3e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 12:14:57 +0900 Subject: [PATCH 2/6] Keep the result set module name on per-cursor dict classes Classes created with type() under ABCMeta took their __module__ from the abc module, so they were shown as abc.AthenaDictResultSet. Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 5 ++++- pyathena/async_cursor.py | 5 ++++- pyathena/cursor.py | 5 ++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index e34e9871c..2d587d8a1 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -262,6 +262,9 @@ def __init__(self, **kwargs) -> None: type( AthenaAioDictResultSet.__name__, (AthenaAioDictResultSet,), - {"dict_type": kwargs["dict_type"]}, + { + "__module__": AthenaAioDictResultSet.__module__, + "dict_type": kwargs["dict_type"], + }, ), ) diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index a0766b484..03187471f 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -351,6 +351,9 @@ def __init__(self, **kwargs) -> None: type( AthenaDictResultSet.__name__, (AthenaDictResultSet,), - {"dict_type": kwargs["dict_type"]}, + { + "__module__": AthenaDictResultSet.__module__, + "dict_type": kwargs["dict_type"], + }, ), ) diff --git a/pyathena/cursor.py b/pyathena/cursor.py index f71302858..e17159dd1 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -232,6 +232,9 @@ def __init__(self, **kwargs) -> None: type( AthenaDictResultSet.__name__, (AthenaDictResultSet,), - {"dict_type": kwargs["dict_type"]}, + { + "__module__": AthenaDictResultSet.__module__, + "dict_type": kwargs["dict_type"], + }, ), ) From b2ece224398082f9601d7928a6f5d8c03aa06c03 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 12:31:31 +0900 Subject: [PATCH 3/6] Apply dict_type to a result set class assigned after initialization A DictCursor subclass that assigned its own AthenaDictResultSet subclass to _result_set_class after DictCursor.__init__ replaced the class that carried dict_type, so its rows were built with dict. The dict cursors now make _result_set_class a property whose setter wraps any assigned AthenaDictResultSet subclass with the cursor's dict_type. An explicit dict_type=None uses the default dict type. Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 41 ++++++++++++++++++++++------- pyathena/async_cursor.py | 40 +++++++++++++++++++++------- pyathena/cursor.py | 40 +++++++++++++++++++++------- tests/pyathena/aio/test_cursor.py | 19 ++++++++++++- tests/pyathena/test_async_cursor.py | 19 +++++++++++-- tests/pyathena/test_cursor.py | 17 +++++++++++- 6 files changed, 145 insertions(+), 31 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 2d587d8a1..bf2b7d719 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -19,6 +19,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions +from pyathena.result_set import AthenaDictResultSet from pyathena.util import override _logger = logging.getLogger(__name__) @@ -254,17 +255,39 @@ def __init__(self, **kwargs) -> None: ``dict_type``, it is the type used to build each row of this cursor's result sets; other cursors are not affected. """ + self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaAioDictResultSet - if "dict_type" in kwargs: - self._result_set_class = cast( - type[AthenaAioDictResultSet], + + @property # type: ignore[explicit-override] # python/mypy#15900 + @override + def _result_set_class(self) -> type[AthenaAioResultSet]: + """The result set class this cursor instantiates for each query. + + Returns: + The class last assigned to this property, or a subclass of it + that carries this cursor's ``dict_type``. + """ + return self._dict_result_set_class + + @_result_set_class.setter + def _result_set_class(self, value: type[AthenaAioResultSet]) -> None: + """Set the result set class this cursor instantiates for each query. + + If this cursor was given ``dict_type`` and ``value`` is a subclass of + ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is + that type is stored instead, so that ``value`` itself is not modified. + + Args: + value: The result set class to instantiate. + """ + if self._dict_type is not None and issubclass(value, AthenaDictResultSet): + value = cast( + type[AthenaAioResultSet], type( - AthenaAioDictResultSet.__name__, - (AthenaAioDictResultSet,), - { - "__module__": AthenaAioDictResultSet.__module__, - "dict_type": kwargs["dict_type"], - }, + value.__name__, + (value,), + {"__module__": value.__module__, "dict_type": self._dict_type}, ), ) + self._dict_result_set_class = value diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 03187471f..69f5dc9f3 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -343,17 +343,39 @@ def __init__(self, **kwargs) -> None: ``dict_type``, it is the type used to build each row of this cursor's result sets; other cursors are not affected. """ + self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - if "dict_type" in kwargs: - self._result_set_class = cast( - type[AthenaDictResultSet], + + @property # type: ignore[explicit-override] # python/mypy#15900 + @override + def _result_set_class(self) -> type[AthenaResultSet]: + """The result set class this cursor instantiates for each query. + + Returns: + The class last assigned to this property, or a subclass of it + that carries this cursor's ``dict_type``. + """ + return self._dict_result_set_class + + @_result_set_class.setter + def _result_set_class(self, value: type[AthenaResultSet]) -> None: + """Set the result set class this cursor instantiates for each query. + + If this cursor was given ``dict_type`` and ``value`` is a subclass of + ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is + that type is stored instead, so that ``value`` itself is not modified. + + Args: + value: The result set class to instantiate. + """ + if self._dict_type is not None and issubclass(value, AthenaDictResultSet): + value = cast( + type[AthenaResultSet], type( - AthenaDictResultSet.__name__, - (AthenaDictResultSet,), - { - "__module__": AthenaDictResultSet.__module__, - "dict_type": kwargs["dict_type"], - }, + value.__name__, + (value,), + {"__module__": value.__module__, "dict_type": self._dict_type}, ), ) + self._dict_result_set_class = value diff --git a/pyathena/cursor.py b/pyathena/cursor.py index e17159dd1..f8a21f29a 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -224,17 +224,39 @@ def __init__(self, **kwargs) -> None: ``dict_type``, it is the type used to build each row of this cursor's result sets; other cursors are not affected. """ + self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - if "dict_type" in kwargs: - self._result_set_class = cast( - type[AthenaDictResultSet], + + @property # type: ignore[explicit-override] # python/mypy#15900 + @override + def _result_set_class(self) -> type[AthenaResultSet]: + """The result set class this cursor instantiates for each query. + + Returns: + The class last assigned to this property, or a subclass of it + that carries this cursor's ``dict_type``. + """ + return self._dict_result_set_class + + @_result_set_class.setter + def _result_set_class(self, value: type[AthenaResultSet]) -> None: + """Set the result set class this cursor instantiates for each query. + + If this cursor was given ``dict_type`` and ``value`` is a subclass of + ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is + that type is stored instead, so that ``value`` itself is not modified. + + Args: + value: The result set class to instantiate. + """ + if self._dict_type is not None and issubclass(value, AthenaDictResultSet): + value = cast( + type[AthenaResultSet], type( - AthenaDictResultSet.__name__, - (AthenaDictResultSet,), - { - "__module__": AthenaDictResultSet.__module__, - "dict_type": kwargs["dict_type"], - }, + value.__name__, + (value,), + {"__module__": value.__module__, "dict_type": self._dict_type}, ), ) + self._dict_result_set_class = value diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 8990a1a99..4561b4613 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -9,7 +9,8 @@ from botocore.exceptions import ClientError from pyathena import BINARY, Binary, ExecuteOptions -from pyathena.aio.cursor import AioCursor +from pyathena.aio.cursor import AioCursor, AioDictCursor +from pyathena.aio.result_set import AthenaAioDictResultSet from pyathena.error import DatabaseError, OperationalError, ProgrammingError from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaQueryExecution @@ -905,3 +906,19 @@ async def test_dict_type(self, aio_dict_cursor): # dict_type of another cursor does not change the row type of this one. await aio_dict_cursor.execute("SELECT * FROM one_row") assert type(await aio_dict_cursor.fetchone()) is dict + + async def test_dict_type_custom_result_set(self, aio_dict_cursor): + class CustomResultSet(AthenaAioDictResultSet): + pass + + class CustomDictCursor(AioDictCursor): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._result_set_class = CustomResultSet + + async with aio_dict_cursor.connection.cursor( + CustomDictCursor, dict_type=OrderedDict + ) as cursor: + await cursor.execute("SELECT * FROM one_row") + assert isinstance(cursor.result_set, CustomResultSet) + assert type(await cursor.fetchone()) is OrderedDict diff --git a/tests/pyathena/test_async_cursor.py b/tests/pyathena/test_async_cursor.py index 710292686..d08ed5801 100644 --- a/tests/pyathena/test_async_cursor.py +++ b/tests/pyathena/test_async_cursor.py @@ -6,10 +6,10 @@ import pytest -from pyathena.async_cursor import AsyncCursor +from pyathena.async_cursor import AsyncCursor, AsyncDictCursor from pyathena.error import NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaResultSet +from pyathena.result_set import AthenaDictResultSet, AthenaResultSet from tests import ENV from tests.pyathena.conftest import connect @@ -255,3 +255,18 @@ def test_dict_type(self, async_dict_cursor): # dict_type of another cursor does not change the row type of this one. _, future = async_dict_cursor.execute("SELECT * FROM one_row") assert type(future.result().fetchone()) is dict + + def test_dict_type_custom_result_set(self, async_dict_cursor): + class CustomResultSet(AthenaDictResultSet): + pass + + class CustomDictCursor(AsyncDictCursor): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._result_set_class = CustomResultSet + + with async_dict_cursor.connection.cursor(CustomDictCursor, dict_type=OrderedDict) as cursor: + _, future = cursor.execute("SELECT * FROM one_row") + result_set = future.result() + assert isinstance(result_set, CustomResultSet) + assert type(result_set.fetchone()) is OrderedDict diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index e3f440e67..22afd5781 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -32,9 +32,10 @@ ) from pyathena.async_cursor import AsyncCursor from pyathena.converter import _to_array, _to_map, _to_struct -from pyathena.cursor import Cursor +from pyathena.cursor import Cursor, DictCursor from pyathena.error import DatabaseError, NotSupportedError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.result_set import AthenaDictResultSet from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -1885,6 +1886,20 @@ def test_dict_type(self, dict_cursor): dict_cursor.execute("SELECT * FROM one_row") assert type(dict_cursor.fetchone()) is dict + def test_dict_type_custom_result_set(self, dict_cursor): + class CustomResultSet(AthenaDictResultSet): + pass + + class CustomDictCursor(DictCursor): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._result_set_class = CustomResultSet + + with dict_cursor.connection.cursor(CustomDictCursor, dict_type=OrderedDict) as cursor: + cursor.execute("SELECT * FROM one_row") + assert isinstance(cursor.result_set, CustomResultSet) + assert type(cursor.fetchone()) is OrderedDict + def test_null_vs_empty_string(self, dict_cursor): """ DictCursor should properly distinguish NULL from empty string. From 2ee5fbc006b51c41a29d4ef9635228bc72bb9c85 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 12:32:31 +0900 Subject: [PATCH 4/6] Say that dict_type=None uses the default row type Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 4 ++-- pyathena/async_cursor.py | 4 ++-- pyathena/cursor.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index bf2b7d719..90bea3de9 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -252,8 +252,8 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``AioCursor.__init__``. If they include - ``dict_type``, it is the type used to build each row of this - cursor's result sets; other cursors are not affected. + ``dict_type`` other than None, it is the type used to build each + row of this cursor's result sets; other cursors are not affected. """ self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 69f5dc9f3..bd30c0f0d 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -340,8 +340,8 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``AsyncCursor.__init__``. If they include - ``dict_type``, it is the type used to build each row of this - cursor's result sets; other cursors are not affected. + ``dict_type`` other than None, it is the type used to build each + row of this cursor's result sets; other cursors are not affected. """ self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) diff --git a/pyathena/cursor.py b/pyathena/cursor.py index f8a21f29a..c72c8b48e 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -221,8 +221,8 @@ def __init__(self, **kwargs) -> None: Args: **kwargs: Arguments forwarded to ``Cursor.__init__``. If they include - ``dict_type``, it is the type used to build each row of this - cursor's result sets; other cursors are not affected. + ``dict_type`` other than None, it is the type used to build each + row of this cursor's result sets; other cursors are not affected. """ self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) From e27e6e206520264a6e5ada11cfed553c50e37134 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 13:02:07 +0900 Subject: [PATCH 5/6] Pass dict_type to the dict result set as a constructor argument Replace the per-cursor result set subclasses and the _result_set_class property with plain argument passing. AthenaDictResultSet takes an optional dict_type that overrides the class attribute for that instance, AthenaAioResultSet.create() forwards extra arguments to the constructor, and the base cursors pass _result_set_kwargs, which the dict cursors set to their dict_type, to each result set they create. Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 46 ++++++-------------------------------- pyathena/aio/result_set.py | 4 ++++ pyathena/async_cursor.py | 45 ++++++------------------------------- pyathena/cursor.py | 45 ++++++------------------------------- pyathena/result_set.py | 13 +++++++++++ 5 files changed, 38 insertions(+), 115 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 90bea3de9..626ff7d2f 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -19,7 +19,6 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.result_set import AthenaDictResultSet from pyathena.util import override _logger = logging.getLogger(__name__) @@ -86,6 +85,7 @@ def __init__( ) self._result_set: AthenaAioResultSet | None = None self._result_set_class = AthenaAioResultSet + self._result_set_kwargs: dict[str, Any] = {} @property # type: ignore[explicit-override] # python/mypy#15900 @override @@ -175,6 +175,7 @@ async def execute( self.arraysize, self._retry_config, result_set_type_hints=options.result_set_type_hints, + **self._result_set_kwargs, ) else: raise OperationalError(query_execution.state_change_reason) @@ -247,47 +248,14 @@ class AioDictCursor(AioCursor): ... print(row["name"]) """ - def __init__(self, **kwargs) -> None: + def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """Initialize an AioDictCursor. Args: - **kwargs: Arguments forwarded to ``AioCursor.__init__``. If they include - ``dict_type`` other than None, it is the type used to build each - row of this cursor's result sets; other cursors are not affected. + dict_type: The type used to build each row of this cursor's result + sets. If None, the result set class's ``dict_type`` is used. + **kwargs: Arguments forwarded to ``AioCursor.__init__``. """ - self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaAioDictResultSet - - @property # type: ignore[explicit-override] # python/mypy#15900 - @override - def _result_set_class(self) -> type[AthenaAioResultSet]: - """The result set class this cursor instantiates for each query. - - Returns: - The class last assigned to this property, or a subclass of it - that carries this cursor's ``dict_type``. - """ - return self._dict_result_set_class - - @_result_set_class.setter - def _result_set_class(self, value: type[AthenaAioResultSet]) -> None: - """Set the result set class this cursor instantiates for each query. - - If this cursor was given ``dict_type`` and ``value`` is a subclass of - ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is - that type is stored instead, so that ``value`` itself is not modified. - - Args: - value: The result set class to instantiate. - """ - if self._dict_type is not None and issubclass(value, AthenaDictResultSet): - value = cast( - type[AthenaAioResultSet], - type( - value.__name__, - (value,), - {"__module__": value.__module__, "dict_type": self._dict_type}, - ), - ) - self._dict_result_set_class = value + self._result_set_kwargs = {"dict_type": dict_type} diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index 5a628da7e..59299f101 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -75,6 +75,7 @@ async def create( arraysize: int, retry_config: RetryConfig, result_set_type_hints: dict[str | int, str] | None = None, + **kwargs: Any, ) -> AthenaAioResultSet: """Async factory method. @@ -88,6 +89,8 @@ async def create( retry_config: Retry configuration for API calls. result_set_type_hints: Optional dictionary mapping column names to Athena DDL type signatures for precise type conversion. + **kwargs: Additional arguments passed to the constructor of ``cls``, + such as ``dict_type`` for ``AthenaAioDictResultSet``. Returns: A fully initialized ``AthenaAioResultSet``. @@ -99,6 +102,7 @@ async def create( arraysize, retry_config, result_set_type_hints=result_set_type_hints, + **kwargs, ) if result_set.state == AthenaQueryExecution.STATE_SUCCEEDED: await result_set._async_pre_fetch() diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index bd30c0f0d..e6ff46427 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -112,6 +112,7 @@ def __init__( self._max_workers = max_workers self._executor = ThreadPoolExecutor(max_workers=max_workers) self._result_set_class = AthenaResultSet + self._result_set_kwargs: dict[str, Any] = {} @property def arraysize(self) -> int: @@ -198,6 +199,7 @@ def _collect_result_set( arraysize=self._arraysize, retry_config=self._retry_config, result_set_type_hints=result_set_type_hints, + **self._result_set_kwargs, ) @override @@ -335,47 +337,14 @@ class AsyncDictCursor(AsyncCursor): >>> print(f"User: {row['name']} ({row['email']})") """ - def __init__(self, **kwargs) -> None: + def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """Initialize an AsyncDictCursor. Args: - **kwargs: Arguments forwarded to ``AsyncCursor.__init__``. If they include - ``dict_type`` other than None, it is the type used to build each - row of this cursor's result sets; other cursors are not affected. + dict_type: The type used to build each row of this cursor's result + sets. If None, the result set class's ``dict_type`` is used. + **kwargs: Arguments forwarded to ``AsyncCursor.__init__``. """ - self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - - @property # type: ignore[explicit-override] # python/mypy#15900 - @override - def _result_set_class(self) -> type[AthenaResultSet]: - """The result set class this cursor instantiates for each query. - - Returns: - The class last assigned to this property, or a subclass of it - that carries this cursor's ``dict_type``. - """ - return self._dict_result_set_class - - @_result_set_class.setter - def _result_set_class(self, value: type[AthenaResultSet]) -> None: - """Set the result set class this cursor instantiates for each query. - - If this cursor was given ``dict_type`` and ``value`` is a subclass of - ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is - that type is stored instead, so that ``value`` itself is not modified. - - Args: - value: The result set class to instantiate. - """ - if self._dict_type is not None and issubclass(value, AthenaDictResultSet): - value = cast( - type[AthenaResultSet], - type( - value.__name__, - (value,), - {"__module__": value.__module__, "dict_type": self._dict_type}, - ), - ) - self._dict_result_set_class = value + self._result_set_kwargs = {"dict_type": dict_type} diff --git a/pyathena/cursor.py b/pyathena/cursor.py index c72c8b48e..fd8d4737c 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -90,6 +90,7 @@ def __init__( **kwargs, ) self._result_set_class = AthenaResultSet + self._result_set_kwargs: dict[str, Any] = {} @property # type: ignore[explicit-override] # python/mypy#15900 @override @@ -192,6 +193,7 @@ def execute( self.arraysize, self._retry_config, result_set_type_hints=options.result_set_type_hints, + **self._result_set_kwargs, ) else: raise OperationalError(query_execution.state_change_reason) @@ -216,47 +218,14 @@ class DictCursor(Cursor): ... print(f"Product {row['id']}: {row['name']} - ${row['price']}") """ - def __init__(self, **kwargs) -> None: + def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """Initialize a DictCursor. Args: - **kwargs: Arguments forwarded to ``Cursor.__init__``. If they include - ``dict_type`` other than None, it is the type used to build each - row of this cursor's result sets; other cursors are not affected. + dict_type: The type used to build each row of this cursor's result + sets. If None, the result set class's ``dict_type`` is used. + **kwargs: Arguments forwarded to ``Cursor.__init__``. """ - self._dict_type: type[Any] | None = kwargs.get("dict_type") super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - - @property # type: ignore[explicit-override] # python/mypy#15900 - @override - def _result_set_class(self) -> type[AthenaResultSet]: - """The result set class this cursor instantiates for each query. - - Returns: - The class last assigned to this property, or a subclass of it - that carries this cursor's ``dict_type``. - """ - return self._dict_result_set_class - - @_result_set_class.setter - def _result_set_class(self, value: type[AthenaResultSet]) -> None: - """Set the result set class this cursor instantiates for each query. - - If this cursor was given ``dict_type`` and ``value`` is a subclass of - ``AthenaDictResultSet``, a subclass of ``value`` whose ``dict_type`` is - that type is stored instead, so that ``value`` itself is not modified. - - Args: - value: The result set class to instantiate. - """ - if self._dict_type is not None and issubclass(value, AthenaDictResultSet): - value = cast( - type[AthenaResultSet], - type( - value.__name__, - (value,), - {"__module__": value.__module__, "dict_type": self._dict_type}, - ), - ) - self._dict_result_set_class = value + self._result_set_kwargs = {"dict_type": dict_type} diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 152715fdb..adbb36c96 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -818,6 +818,19 @@ class AthenaDictResultSet(AthenaResultSet): # You can override this to use OrderedDict or other dict-like types. dict_type: type[Any] = dict + def __init__(self, *args: Any, dict_type: type[Any] | None = None, **kwargs: Any) -> None: + """Initialize the result set with an optional row type for this instance. + + Args: + *args: Positional arguments passed to the next ``__init__`` in the MRO. + dict_type: The type used to build each row of this result set. If + None, the class attribute ``dict_type`` is used. + **kwargs: Keyword arguments passed to the next ``__init__`` in the MRO. + """ + if dict_type is not None: + self.dict_type = dict_type + super().__init__(*args, **kwargs) + @override def _get_rows( self, From ee872f7a91da39ff2e60969718c57366a723ee6e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 3 Oct 2026 13:04:29 +0900 Subject: [PATCH 6/6] Pass dict_type to the result set only when it is given A result set subclass whose __init__ does not accept dict_type keeps working with dict cursors created without dict_type. Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/cursor.py | 3 ++- pyathena/async_cursor.py | 3 ++- pyathena/cursor.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 626ff7d2f..00a57be0b 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -258,4 +258,5 @@ def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """ super().__init__(**kwargs) self._result_set_class = AthenaAioDictResultSet - self._result_set_kwargs = {"dict_type": dict_type} + if dict_type is not None: + self._result_set_kwargs = {"dict_type": dict_type} diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index e6ff46427..2963e5f40 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -347,4 +347,5 @@ def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """ super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - self._result_set_kwargs = {"dict_type": dict_type} + if dict_type is not None: + self._result_set_kwargs = {"dict_type": dict_type} diff --git a/pyathena/cursor.py b/pyathena/cursor.py index fd8d4737c..099f8e8eb 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -228,4 +228,5 @@ def __init__(self, dict_type: type[Any] | None = None, **kwargs) -> None: """ super().__init__(**kwargs) self._result_set_class = AthenaDictResultSet - self._result_set_kwargs = {"dict_type": dict_type} + if dict_type is not None: + self._result_set_kwargs = {"dict_type": dict_type}