Add typings #577

bryanforbes · 2020-05-18T19:09:57Z

This is a work in progress and still produces errors when type checking the codebase. I added type hints to the Python files where it was feasible and updated the code of the methods to be type safe. In two places (pool.py and exceptions/_base.py), adding type hints to the code itself would not work because of the dynamic nature of the code in those modules. There may also be places where I'm making wrong assumptions about the code (especially in the cython code), so a thorough review would be very welcome. Lastly, I did not add typings for asyncpg._testbase since that seems to be an internal testing module.

Some of the remaining problems that mypy finds may be bugs in the code or they may be that mypy is being overly strict:

cursor.py: In BaseCursor.__repr__, self._state could technically be None, and cause an exception when self._state.query is used
connection.py:
- In Python 3.7+, asyncio.current_task() can return None, so compat.current_task() has to be typed as returning Optional[Task[Any]]. This means Connection._cancel() may throw an exception when self._cancellations.discard(compat.current_task(self._loop)) is called
- The code in _extract_stack() has a couple of issues:
  - Passing an iterator to StackSummary.extract() but it expects a generator
  - __path__ does not exist on the asyncpg module
connect_utils.py has several errors that relate to the code in _parse_connect_dsn_and_args() assigning new types to variables originally declared as another type. I can try to clean this up to make it more type-safe, or just add # type: ignore.
cluster.py:
- There are a couple of places where a bytes is being passed to a formatting string which mypy recommends using !r with those
- In Cluster.connect(), self.get_connection_spec() can return None, which would cause conn_info.update() to throw an error. Is this desired?
- A similar issue can be found in Cluster._test_connection() with self._connection_addr

References #569, #387

bryanforbes · 2020-05-19T00:32:16Z

One thing to note is that the handling of Record isn't quite what I'd like it to be yet. Currently, a user could do something like this:

class MyRecord(asyncpg.Record):
    id: int
    name: str
    created_at: datetime.datetime

r: MyRecord = ...
reveal_type(r.id) # int

The user will only use this class for type checking and can use attribute access with type checking, but not lookup notation. A mypy plugin can probably be written to handle the lookup notation (TypedDict has something similar), so this isn't a big issue. The bigger issue is the following will be an error:

records: typing.List[MyRecord] = await conn.fetch('SELECT ... from ...')

It's an error because fetch() is typed to return List[Record], which doesn't match List[MyRecord] (List is invariant, so it has to match exactly). There are two options:

Return a bound generic from every function that returns a Record. This mostly works, but it requires the user to type their variables, even if they're not using a custom class:

record: MyRecord = await conn.fetchrow('SELECT ... FROM ...')
records: typing.List[MyRecord] = await conn.fetch('SELECT ... FROM ...')
cursor: asyncpg.cursor.CursorFactory[MyRecord] = await conn.cursor('SELECT ... FROM ...')
record2: asyncpg.Record = await conn.fetchrow('SELECT ... FROM ...')
records2: typing.List[asyncpg.Record] = await conn.fetch('SELECT ... FROM ...')
cursor2: asyncpg.cursor.CursorFactory[asyncpg.Record] = await conn.cursor('SELECT ... FROM ...')

Add a record_class parameter to methods to any function that returns a Record. This, in my opinion, works the best but adds an unused (code-wise) parameter:

# the variables below have the same types as the last code block,
# but the types are inferred instead
record = await conn.fetch('SELECT ... FROM ...', record_class=MyRecord)
records = await conn.fetch('SELECT ... FROM ...', record_class=MyRecord)
cursor = await conn.cursor('SELECT ... FROM ...', record_class=MyRecord)
record2 = await conn.fetchrow('SELECT ... FROM ...')
records2 = await conn.fetch('SELECT ... FROM ...')
cursor2 = await conn.cursor('SELECT ... FROM ...')

Note that while mypy will see record as a MyRecord, it will be an instance of Record at runtime because record_class is only used to infer the type of the return.

I have the second option coded locally, but I wanted to check if adding that extra record_class parameter for type-checking purposes is OK.

elprans · 2020-05-19T02:54:11Z

There's a third option: lie about the return type and say it's a Sequence[Record] instead. I don't think the results of fetch() are mutated that often (why would you?) so this shouldn't be much of an issue.

bryanforbes · 2020-05-19T03:24:49Z

Typing it as Sequence[Record] kind of works, but you'd still need a cast to go from Sequence[Record] to Sequence[MyRecord]. Combining Sequence with option 2 is better than List (since Sequence is covariant), but it still requires adding type information for all of your variables storing the results from functions returning Record.

elprans · 2020-05-19T04:41:59Z

but you'd still need a cast

Right. Well, my concern with the record_class parameter is that it could be very confusing to unknowing users. If we named it record_class, then one might rightfully assume that the returned records will be instances of that class, when, in fact, they wouldn't be.

You'd mentioned that a plugin is necessary to handle subscript access, so perhaps said plugin can also deal with the fetch calls (via a function hook). The plugin would only need to check that the lval type is a covariant sequence of Record and adjust the return type of fetch() to it to make mypy happy.

bryanforbes · 2020-05-19T13:40:29Z

Right. Well, my concern with the record_class parameter is that it could be very confusing to unknowing users.

I understand the concern and share your concern. Would record_type be less confusing?

The plugin would only need to check that the lval type is a covariant sequence of Record and adjust the return type of fetch() to it to make mypy happy.

I'll have to do some digging to see if this is possible. If it is, I agree that it would probably be the best solution.

bryanforbes · 2020-06-04T18:40:30Z

I've done quite a bit more work on the mypy plugin and have the following working:

import asyncpg
import datetime
import typing
import typing_extensions

class MyRecord(asyncpg.Record):
    foo: int
    bar: typing.Optional[str]

class MyOtherRecord(MyRecord):
    baz: datetime.datetime

async def main() -> None:
    conn = await asyncpg.connect(...)
    m = typing.cast(MyRecord, await conn.fetchrow('SELECT foo, bar FROM records'))
    o = typing.cast(MyOtherRecord, await conn.fetchrow('SELECT foo, bar, baz FROM other_records'))

    key = 'baz'
    fkey: typing_extensions.Final = 'baz'

    reveal_type(m['foo'])  # int
    reveal_type(m['bar'])  # Optional[str]
    reveal_type(m['baz'])  # error: "MyRecord" has no key 'baz'
    reveal_type(m[0])  # int
    reveal_type(m[1])  # Optional[str]
    reveal_type(m[2])  # error: "MyRecord" has no index 2
    reveal_type(m.get('foo'))  # int
    reveal_type(m.get('bar'))  # Optional[str]
    reveal_type(m.get('baz'))  # error: "MyRecord" has no key 'baz'
    reveal_type(m.get('baz', 1))  # Literal[1]
    reveal_type(m.get(key, 1))  # Union[Any, int]
    reveal_type(m.get(fkey, 1))  # Literal[1]

    reveal_type(o['foo'])  # int
    reveal_type(o['bar'])  # Optional[str]
    reveal_type(o['baz'])  # datetime
    reveal_type(o[0])  # int
    reveal_type(o[1])  # Optional[str]
    reveal_type(o[2])  # datetime
    reveal_type(o.get('foo'))  # int
    reveal_type(o.get('bar'))  # Optional[str]
    reveal_type(o.get('baz'))  # datetime
    reveal_type(o.get('baz', 1))  # datetime
    reveal_type(o.get(key, 1))  # Union[Any, int]
    reveal_type(o.get(fkey, 1))  # datetime

Based on the implementation of asyncpg.Record, I believe I have the typing for __getitem__() and get() correct. I tried to get the typings for Record to be as similar to TypedDict as possible (given the implementation differences). You'll notice that when the key can be determined by the type system, get() with a default argument is deterministic (ex. o.get('baz', 1) and o.get(fkey, 1)), otherwise it returns a Union. One thing I'd like to possibly try is to come up with a metaclass that would act like a typing_extensions.Protocol with runtime checking so instanceof() could be used to determine if a Record matched. At this time, I haven't attempted it.

I also did a lot of poking around to try and get the plugin to set the return type based on the variable it is being assigned to, but I don't see a way that it's possible. This means we're left with two options:

Force users to cast to subclasses of asyncpg.Record (as seen in the examples above). This means types like PreparedStatement would need to be exposed from asyncpg so users could easily cast the result of Connection.prepare() to asyncpg.PreparedStatement[MyRecord]:

stmt = typing.cast(asyncpg.prepared_stmt.PreparedStatement[MyRecord], await conn.prepare('SELECT ... FROM ...'))
reveal_type(stmt)  # asyncpg.prepared_stmt.PreparedStatement[MyRecord]

Add an unused parameter to indicate to the type system which type should be inferred (and if it's left off, the return type would be asyncpg.Record). I suggested record_class above, but I think return_type would be less prone to users thinking the result would be a different class. This approach would mean that the result of calls to functions like prepare() wouldn't need to be cast to a subscript and the type system would infer the subscript:

stmt = await conn.prepare('SELECT ... FROM ...', return_type=MyRecord)
reveal_type(stmt)  # asyncpg.prepared_stmt.PreparedStatement[MyRecord]

There's also a possible third option if the Protocol-like class actually works: require users to use isinstance() to narrow the type:

stmt = await conn.prepare('SELECT ... FROM ...')
reveal_type(stmt)  # asyncpg.prepared_stmt.PreparedStatement[asyncpg.protocol.protocol.Record]

record = await stmt.fetchrow(...)
reveal_type(record)  # Optional[asyncpg.protocol.protocol.Record]
assert isinstance(record, MyRecord)
reveal_type(record)  # MyRecord

The third option completely depends on whether the Protocol-like class is feasible, and would also do runtime checking (which would check if Record.keys() has all of the keys of the subclass). I would imagine the runtime checking would be a performance hit.

elprans · 2020-06-05T18:03:59Z

Awesome work on the plugin, @bryanforbes! Thanks!

Add an unused parameter to indicate to the type system which type should be inferred

I'm still a bit uneasy with adding unused parameters just for the typing purpose. That said, if we had record_class= actually make fetch() and friends return instances of that class, that would be a great solution. There's actually #40, which quite a few people requested before.

bryanforbes · 2020-06-08T14:31:19Z

@elprans Thanks! There are still some places in the code where # type: ignore is being used. I've updated them to use the specific code(s) that can be used to ignore them (so if new type errors arise, mypy doesn't ignore those as well). Some of them can be ignored. For example, anywhere TypedDict is being used to ensure a dictionary has the right shape, mypy will complain when using ** on it. Others indicate a possible issue with the code that I wasn't sure how to fix. I'll list those here (using the latest commit I've just pushed):

connection.py line 1264: All of the methods use _protocol without checking if it's None, so I typed it as BaseProtocol. However, when the connection is aborted None is assigned to _protocol so _protocol should technically be typed as Optional[BaseProtocol] and be checked everywhere it's used (with a descriptive error) so the methods of the Connection don't throw a strange error about an internal member if they're used after aborting the connection.
connection.py line 1357: compat.current_asyncio_task() can return None (because the standard library can), _cancellations is typed as Set[asyncio.Task[Any]], so discard() rejects passing None to it.
cluster.py lines 129, 547, 604: mypy gives the following error: On Python 3 '{}'.format(b'abc') produces "b'abc'"; use !r if this is a desired behavior. This probably should be updated to use !r, but I wasn't sure.
cluster.py line 136: get_connection_spec() can return None which would throw an error about an internal variable being None. This should probably be checked to see if it's None.
cursor.py line 169: _state can be None, which would cause an exception to be raised here.

With regards to the unused parameter, I completely understand your concern. I also think that making fetch() and friends return the instances of that class would be a good solution. I can look into that, but my experience with cpython is fairly limited (I've used cython a long long time ago, so it's much more familiar). Let me know how you'd like to proceed.

Add the new `record_class` parameter to the `create_pool()` and `connect()` functions, as well as to the `cursor()`, `prepare()`, `fetch()` and `fetchrow()` connection methods. This not only allows adding custom functionality to the returned objects, but also assists with typing (see #577 for discussion). Fixes: #40.

elprans

With regards to the unused parameter, I completely understand your concern. I also think that making fetch() and friends return the instances of that class would be a good solution.

I took this upon myself to implement in #559.

I did a cursory review of the PR and left a few comments. Most importantly, I think we should bite the bullet and use the PEP 526 syntax for type annotations instead of comments. Python 3.5 is rapidly going out of fashion, and I'd love to drop a bunch of hacks we already have to support it.

Thanks again for working on this!

elprans · 2020-07-18T16:43:55Z

asyncpg/cluster.py

    low, high = port_range

-    port = low
+    port = low  # type: typing.Optional[int]


Python 3.5 would be EOL'd in September and I think we should just drop support for it and use proper annotation syntax everywhere.

I've updated the PR to switch to 3.6 type annotations and removed 3.5 from CI

asyncpg/cluster.py

asyncpg/connect_utils.py

elprans · 2020-07-18T16:55:49Z

asyncpg/connection.py

        # Put the connection into the aborted state.
        self._aborted = True
        self._protocol.abort()
-        self._protocol = None
+        self._protocol = None  # type: ignore[assignment]


Perhaps a cleaner solution would be to assign a sentinel instance of a Protocol-like object, e.g. DeadProtocol that raises an error on any attribute access.

Add the new `record_class` parameter to the `create_pool()` and `connect()` functions, as well as to the `cursor()`, `prepare()`, `fetch()` and `fetchrow()` connection methods. This not only allows adding custom functionality to the returned objects, but also assists with typing (see #577 for discussion). Fixes: #40.

victoraugustolls · 2020-08-15T17:12:50Z

@bryanforbes PR #599 was merged! I don't know if it was a blocker for this PR or not (by reading the discussion here I think it was). Thanks for this PR, really, looking forward to it! 😄

bryanforbes · 2020-08-20T15:19:20Z

@victoraugustolls I'll rebase and update this PR today

bryanforbes · 2020-08-21T20:06:08Z

@victoraugustolls I finished the rebase, but there's an issue with Python 3.6 and ConnectionMeta + Generic related to python/typing#449. I'll poke around with it this weekend and try to get it working.

victoraugustolls · 2020-08-22T19:16:40Z

Thanks for the update @bryanforbes ! I will try and search for something that can help

bryanforbes · 2020-08-24T14:04:09Z

@victoraugustolls I was able to come up with a solution for the metaclass issue in Python 3.6 (that won't affect performance on 3.7+), but I'd like to work on some tests before taking this PR out of draft. Feel free to review the code, though.

bryanforbes · 2020-10-16T17:26:22Z

@victoraugustolls I started writing unit tests last night, however they use py.test and pytest-mypy-plugins to allow testing the output for both errors and revealed types. However, although py.test is a dev extra, I noticed you aren't using it to run the test suite. Will it be a problem to use py.test to only run the mypy unit tests?

elprans · 2020-10-16T17:59:44Z

@bryanforbes I don't think we need more than running mypy like we do flake8 in test__sourcecode.py for the main regression testsuite. Here's an example

victoraugustolls · 2020-10-16T18:07:13Z

Hi @bryanforbes ! I'm not a maintainer, but thanks for asking! I agree with @elprans and running mypy should be enough. If mypy is happy it should be fine!

bryanforbes · 2022-02-08T23:37:36Z

I spent some time today rebasing and cleaning up this PR. I'd appreciate another review if someone can spare the cycles.

HarrySky

Left a few comments/questions, didn't get through all the changes, yet (that's a lot of work you have done, thank you!)

asyncpg/compat.py

HarrySky · 2022-02-09T15:01:27Z

asyncpg/connect_utils.py

+def _read_password_file(
+    passfile: pathlib.Path
+) -> typing.List[typing.Tuple[str, ...]]:


Why is this style not used everywhere? Is line breaks with \ better?

There seems to be a mix of styles in the codebase. I use the black style (seen here) in my own projects, and I saw it in some of the newer code, so I went with it. If that needs to be changed, I can format the code a different way.

I prefer Black style, but it is up to asyncpg maintainers

Same here, but I'll wait for the maintainers 😄

asyncpg/connect_utils.py

asyncpg/connection.py

asyncpg/introspection.py

bryanforbes · 2022-02-09T15:55:21Z

Left a few comments/questions, didn't get through all the changes, yet (that's a lot of work you have done, thank you!)

Thank you! I think I responded or 👍 all your comments. I'll address what I can.

bitterteriyaki · 2022-05-05T03:40:38Z

Bumping this PR. Is there something pending?

bryanforbes · 2022-07-26T16:24:48Z

@elprans could you approve the workflow so the regression tests can run?

Andarius · 2022-12-02T13:29:32Z

Hey guys, thanks for the amazing work.
Do you need help on this PR to get this merged? @elprans @bryanforbes

bryanforbes · 2022-12-02T14:57:00Z

I need to rebase on master. I have some free time coming up in the next few weeks, so I should be able to take care of it. I'll let you know if I need some help, @Andarius.

ethanleifer · 2023-05-22T22:07:21Z

Is there anything blocking this pr? I would be willing to help out if needed

takeda · 2023-05-23T00:16:48Z

Yeah I think it should be merged. Maybe it is still not perfect, but it likely will work for majority of people and any bugs will come out when people start using.

It's been 3 years this PR has been open and rebasing it constantly just adds unnecessary overhead.

ethanleifer · 2023-05-23T00:52:28Z

Is there someone we have to ping to get this merged? I would be willing to help out wherever is needed

takeda · 2023-05-23T04:49:30Z

Is there someone we have to ping to get this merged? I would be willing to help out wherever is needed

I saw @elprans initially commenting here with suggestions, and @1st1 in another issue commenting that type annotations would be welcomed. Hopefully that could bring attention to this PR.

Voldemat · 2023-09-14T21:35:13Z

Hi! When you plan to merge this proposal to main?

r614 · 2024-02-06T22:42:09Z

bumping this thread since its been a while, any ETA on the merge @elprans @1st1 - seems like this PR is basically done?

eirnym · 2024-02-12T20:39:05Z

Hi @bryanforbes, could you please resolve conflicts. I have no permissions to do that

@elprans I see that this work could be continued even after this PR has been merged.

bryanforbes · 2024-02-12T20:44:53Z

I'll try to get some time to resolve conflicts this week

Add typings to the project and check the project using mypy. `PostgresMessage` and `PoolConnectionProxy` were broken out into their own files to make it easier to add typing via stub (pyi) files. Since they are metaclasses which generate dynamic objects, we can't type them directly in their python module.

bryanforbes · 2024-02-15T00:11:50Z

@elprans I've updated the typings to match what I have in asyncpg-stubs. Please give this another review when you have time.

…ntime dep

bryanforbes · 2024-03-04T19:49:18Z

@elprans would it help you review this to break this into individual PRs? There's a lot here.

elprans · 2024-03-04T20:01:29Z

@elprans would it help you review this to break this into individual PRs? There's a lot here.

Yes, absolutely.

bryanforbes · 2024-03-04T20:04:40Z

@elprans would it help you review this to break this into individual PRs? There's a lot here.

Yes, absolutely.

Ok. I'll take care of that this week.

bryanforbes marked this pull request as draft May 26, 2020 14:42

elprans mentioned this pull request Jul 18, 2020

Make PoolAcquireContext return a Connection type #578

Open

elprans mentioned this pull request Jul 19, 2020

Allow using custom Record class #599

Merged

elprans reviewed Jul 19, 2020

View reviewed changes

bryanforbes force-pushed the feature/add-typings branch from d449a6f to da85005 Compare July 19, 2020 22:55

bryanforbes force-pushed the feature/add-typings branch from 4fef52f to 8332c03 Compare August 21, 2020 19:40

bryanforbes force-pushed the feature/add-typings branch from 18f3e1b to 8ac22c1 Compare October 13, 2020 14:05

HarrySky reviewed Feb 9, 2022

View reviewed changes

bryanforbes mentioned this pull request Feb 9, 2022

Improve type completeness MagicStack/py-pgproto#15

Merged

ps-george mentioned this pull request Jul 25, 2022

Support for asyncpg 0.26.0? bryanforbes/asyncpg-stubs#109

Closed

bryanforbes force-pushed the feature/add-typings branch from 23154c1 to 96c0a4c Compare July 26, 2022 16:21

elprans mentioned this pull request Sep 28, 2022

Add ability to set the Record factory on the connection #40

Closed

bryanforbes force-pushed the feature/add-typings branch from fb6d8ef to 348acbb Compare February 14, 2024 23:50

bryanforbes force-pushed the feature/add-typings branch from 348acbb to 276bb46 Compare February 15, 2024 00:08

Use TypeAlias annotation and ensure typing_extensions is not a ru…

5ede989

…ntime dep

sanderr mentioned this pull request Mar 1, 2024

Performance investigation inmanta/inmanta-core#7278

Closed

14 tasks

bryanforbes added 2 commits March 1, 2024 11:03

Update pgproto

d7e8ae7

Add imports for back-compat

4e2c87b

Add typings #577

Are you sure you want to change the base?

Add typings #577

Conversation

bryanforbes commented May 18, 2020 • edited

bryanforbes commented May 19, 2020 • edited

elprans commented May 19, 2020

bryanforbes commented May 19, 2020

elprans commented May 19, 2020

bryanforbes commented May 19, 2020

bryanforbes commented Jun 4, 2020 • edited

elprans commented Jun 5, 2020

bryanforbes commented Jun 8, 2020

elprans left a comment

Choose a reason for hiding this comment

elprans Jul 18, 2020

Choose a reason for hiding this comment

bryanforbes Jul 19, 2020

Choose a reason for hiding this comment

elprans Jul 18, 2020

Choose a reason for hiding this comment

victoraugustolls commented Aug 15, 2020

bryanforbes commented Aug 20, 2020

bryanforbes commented Aug 21, 2020

victoraugustolls commented Aug 22, 2020

bryanforbes commented Aug 24, 2020

bryanforbes commented Oct 16, 2020

elprans commented Oct 16, 2020 • edited

victoraugustolls commented Oct 16, 2020

bryanforbes commented Feb 8, 2022

HarrySky left a comment

Choose a reason for hiding this comment

HarrySky Feb 9, 2022

Choose a reason for hiding this comment

bryanforbes Feb 9, 2022

Choose a reason for hiding this comment

HarrySky Feb 9, 2022

Choose a reason for hiding this comment

bryanforbes Feb 9, 2022

Choose a reason for hiding this comment

bryanforbes commented Feb 9, 2022

bitterteriyaki commented May 5, 2022 • edited

bryanforbes commented Jul 26, 2022

Andarius commented Dec 2, 2022

bryanforbes commented Dec 2, 2022

ethanleifer commented May 22, 2023

takeda commented May 23, 2023

ethanleifer commented May 23, 2023

takeda commented May 23, 2023

Voldemat commented Sep 14, 2023

r614 commented Feb 6, 2024

eirnym commented Feb 12, 2024

bryanforbes commented Feb 12, 2024

bryanforbes commented Feb 15, 2024

bryanforbes commented Mar 4, 2024

elprans commented Mar 4, 2024

bryanforbes commented Mar 4, 2024

bryanforbes commented May 18, 2020 •

edited

bryanforbes commented May 19, 2020 •

edited

bryanforbes commented Jun 4, 2020 •

edited

elprans commented Oct 16, 2020 •

edited

bitterteriyaki commented May 5, 2022 •

edited