Enhances Cursor type inference capabilities

[woosuk-choi-g]

Related Issue

• 484

Overview

Enhances type inference capabilities. While there may be no significant impact in some use cases, it is better compared to before.

Use Case

connection = Connection() # -> pyathena.connection.Connection[pyathena.cursor.Cursor]
connection_cursor = connection.cursor() # -> pyathena.cursor.Cursor

async_connection = Connection(cursor_class=AsyncCursor) # -> pyathena.connection.Connection[pyathena.async_cursor.AsyncCursor]
async_connection_cursor = async_connection.cursor() # -> pyathena.async_cursor.AsyncCursor
dict_cursor = async_connection.cursor(DictCursor) # -> (pyathena.async_cursor.AsyncCursor | pyathena.cursor.DictCursor) (IDE seems to be pretty confusing)

factory_connection = connect(cursor_class=AsyncCursor) # -> pyathena.connection.Connection (factory can't be inferred)
factory_connection_cursor = factory_connection.cursor() # -> Any (factory can't be inferred)
factory_connection_functional_cursor = factory_connection.cursor(DictCursor) # -> pyathena.cursor.DictCursor

Limitation

• using a mix of multiple cursor is inferred a union of cursor types
• def connection factory method still not support type hint

[laughingman7743]

Thank you for your contribution. I have made some comments, please check them out.

[woosuk-choi-g]

Now, the type analyzer can accurately infer the Cursor type of Connection.

# reveal_type.py
connection_aysnc_cursor = connect(cursor_class=AsyncCursor)
reveal_type(connection_aysnc_cursor)

aysnc_cursor = connection_aysnc_cursor.cursor()
reveal_type(aysnc_cursor)

dict_cursor = connection_aysnc_cursor.cursor(DictCursor)
reveal_type(dict_cursor)

connection_cursor = Connection()
reveal_type(connection_cursor)

cursor = connection_cursor.cursor()
reveal_type(cursor)

connection_async = Connection(cursor_class=AsyncCursor)
reveal_type(connection_async)

async_cursor = connection_async.cursor()
reveal_type(async_cursor)

dict_cursor = connection_async.cursor(DictCursor)
reveal_type(dict_cursor)

$ mypy reveal_type.py
.vscode\main.py:8: note: Revealed type is "pyathena.connection.Connection[pyathena.async_cursor.AsyncCursor]"
.vscode\main.py:11: note: Revealed type is "pyathena.async_cursor.AsyncCursor"
.vscode\main.py:14: note: Revealed type is "pyathena.cursor.DictCursor"
.vscode\main.py:17: note: Revealed type is "pyathena.connection.Connection[pyathena.cursor.Cursor]" 
.vscode\main.py:20: note: Revealed type is "pyathena.cursor.Cursor"
.vscode\main.py:23: note: Revealed type is "pyathena.connection.Connection[pyathena.async_cursor.AsyncCursor]"
.vscode\main.py:26: note: Revealed type is "pyathena.async_cursor.AsyncCursor"
.vscode\main.py:29: note: Revealed type is "pyathena.cursor.DictCursor"

and fixed formatting error 😊

$ mypy .
Success: no issues found in 40 source files

[laughingman7743]

I just released v3.2.0. 🎉 https://pypi.org/project/PyAthena/3.2.0/ https://github.com/laughingman7743/PyAthena/releases/tag/v3.2.0

[laughingman7743]

@woosuk-choi-g Please check the following issues: https://github.com/laughingman7743/PyAthena/issues/506