MagicStack / asyncpg

A fast PostgreSQL Database Client Library for Python/asyncio.
Apache License 2.0
6.99k stars 404 forks source link

Typing: Correctly set types when getting `Record` fields #1210

Open galah92 opened 1 week ago

galah92 commented 1 week ago

Consider the following:

import asyncio

import asyncpg

async def main() -> int:
    connection_string = "postgresql://postgres:postgres@127.0.0.1:5432/postgres"
    conn = await asyncpg.connect(connection_string)
    result = await conn.fetch("SELECT 1 + 1 AS num", record_class=Row)
    row = result[0]
    num = row["num"]
    print(type(num))
    return num  # mypy error: Returning Any from function declared to return "int"

class Row(asyncpg.Record):
    num: int

if __name__ == "__main__":
    asyncio.run(main())

In essence, the type of num is Any instead of int. I tried to make Row inherit from TypedDict or use a dataclass, both failed. I'll also be happy to help develop this feature with some general guidance.