Add generic typing support for Memory[Send/Receive]Channel

[jakkdl]

Fixes #1327

I didn't figure out a good way to add tests for it, but I reproduced the error from the original PR, and it now works as expected (in a somewhat pared down variant of it):

from trio import open_memory_channel, MemorySendChannel, MemoryReceiveChannel
from typing import cast, Tuple

OutboundBroadcastChannelPair = Tuple[
    MemorySendChannel[int],
    MemoryReceiveChannel[int],
]

(outbound_send_channel, outbound_receive_channel) = cast(
    OutboundBroadcastChannelPair, open_memory_channel(100)
)

outbound_send_channel.send_nowait(5)
outbound_send_channel.send_nowait(None) # argument 1 has incompatible type

reveal_type(outbound_receive_channel.receive_nowait()) # revealed type is builtins.int

One way would be to add type hints to the tests in trio/tests/test_channel.py, and also check the trio.tests submodule with mypy in check.sh

As noted in https://github.com/python-trio/trio/blob/master/trio/_util.py#L251-L254 @generic_function won't type-check, and I didn't figure out what mypy plugin / clever stub would enable that.

I didn't manage to figure out what trio.lowlevel.wait_task_rescheduled actually returns, so it's possible that the return type on MemoryReceiveChannel.receive is incorrect.

I assumed that the co/contravariantness of the TypeVars in trio/_abc.py for SendChannel and MemoryChannel would apply for Memory[Send/Receive]Channel as well.

While on it I also typed ~everything in trio/_channel.py. (though the __exit__ ones look kinda crazy 😅 )

[codecov]

Codecov Report

Merging #2549 (d5fd09f) into master (ca850f6) will increase coverage by 3.19%.
The diff coverage is 100.00%.

❗ Current head d5fd09f differs from pull request most recent head 798e87a. Consider uploading reports for the commit 798e87a to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2549      +/-   ##
==========================================
+ Coverage   89.26%   92.45%   +3.19%     
==========================================
  Files         118      118              
  Lines       16334    16346      +12     
  Branches     1800     3155    +1355     
==========================================
+ Hits        14580    15113     +533     
+ Misses       1462     1104     -358     
+ Partials      292      129     -163     

Impacted Files	Coverage Δ
trio/_core/__init__.py	75.00% <ø> (ø)
trio/lowlevel.py	58.33% <ø> (ø)
trio/_channel.py	100.00% <100.00%> (+7.26%)	⬆️
trio/_core/_traps.py	100.00% <100.00%> (+8.57%)	⬆️
trio/__init__.py	100.00% <0.00%> (ø)
trio/_windows_pipes.py	0.00% <0.00%> (ø)
trio/_core/_io_windows.py	0.00% <0.00%> (ø)
trio/tests/test_windows_pipes.py	26.58% <0.00%> (ø)
trio/testing/_check_streams.py	100.00% <0.00%> (+0.33%)	⬆️
trio/tests/test_highlevel_open_tcp_stream.py	98.16% <0.00%> (+0.36%)	⬆️
... and 65 more

[A5rocks]

I've got a couple nitpicks re: types

[A5rocks]

cc @agronholm as you've had some opinions on types before (hopefully I'm not confusing you with someone else? My memory is fuzzy...)

[agronholm]

The previous typing PR used from __future__ import annotations, so I see no reason to not do that here too. That lets you modernize the annotations by not using deprecated types from the typing module.

[jakkdl]

Pushed a commit with updates from the reviews, though I got stuck trying to resolve open_memory_channel. It currently complains

trio/_channel.py:100: error: Incompatible return value type
(got "Tuple[MemorySendChannel[SendType], MemoryReceiveChannel[ReceiveType]]",
expected "Tuple[MemorySendChannel[T], MemoryReceiveChannel[T]]")  [return-value]

and I'm in general quite confused how to write that the function, with the @generic_function decorator it's expected to take one type, which is passed to MemoryChannelState, and then we're to return the channel classes typed with two other typevars, presumably bound to the type parameter passed, but with different covariant/contravariantness. So if anybody has an idea how to cram all that into the function header (or class header, see below), I would love to be enlightened.

I looked at trio-typing, and they resort to making open_memory_channel a class:

# written as a class so you can say open_memory_channel[int](5)
class open_memory_channel(Tuple[MemorySendChannel[T], MemoryReceiveChannel[T]]):
    def __new__(  # type: ignore[misc]  # "must return a subtype"
        cls, max_buffer_size: float
    ) -> Tuple[MemorySendChannel[T], MemoryReceiveChannel[T]]: ...
    def __init__(self, max_buffer_size: float): ...

which looks kinda hackish, but maybe is fine? (though this doesn't resolve my typing headaches, and mostly just makes it so you don't need to use @generic_function and can get typechecking working on the "function" without any magic. (And they don't appear to use @generic_function anywhere in trio-typing.))

[jakkdl]

CC @oremanj who wrote the original @generic_function in #908, if you wouldn't mind having opinions here :)

[A5rocks]

Pushed a commit with updates from the reviews, though I got stuck trying to resolve open_memory_channel. It currently complains

trio/_channel.py:100: error: Incompatible return value type
(got "Tuple[MemorySendChannel[SendType], MemoryReceiveChannel[ReceiveType]]",
expected "Tuple[MemorySendChannel[T], MemoryReceiveChannel[T]]")  [return-value]

and I'm in general quite confused how to write that the function, with the @generic_function decorator it's expected to take one type, which is passed to MemoryChannelState, and then we're to return the channel classes typed with two other typevars, presumably bound to the type parameter passed, but with different covariant/contravariantness. So if anybody has an idea how to cram all that into the function header (or class header, see below), I would love to be enlightened.

I think this is because of NoPublicConstructor's _create function. Not sure how to fix it, but hopefully that makes some stuff clearer. Maybe we can override it with a custom _create impl? (That just calls into super())

[jakkdl]

Ah, that made me think of the super simple solution!

    return (
        MemorySendChannel[T]._create(state),
        MemoryReceiveChannel[T]._create(state),
    )

I should definitely write some test cases to check that stuff actually works though.

[agronholm]

Ah, that made me think of the super simple solution!

    return (
        MemorySendChannel[T]._create(state),
        MemoryReceiveChannel[T]._create(state),
    )

I should definitely write some test cases to check that stuff actually works though.

This requires Python 3.8 at minimum, just FYI.

[jakkdl]

Ah, that made me think of the super simple solution!

    return (
        MemorySendChannel[T]._create(state),
        MemoryReceiveChannel[T]._create(state),
    )

I should definitely write some test cases to check that stuff actually works though.

This requires Python 3.8 at minimum, just FYI.

ah, shoot

[agronholm]

Ah, that made me think of the super simple solution!

    return (
        MemorySendChannel[T]._create(state),
        MemoryReceiveChannel[T]._create(state),
    )

I should definitely write some test cases to check that stuff actually works though.

This requires Python 3.8 at minimum, just FYI.

ah, shoot

Actually, no. The class level __getitem__() method was in fact introduced in Python 3.7. It's just the collections.abc module that gained such support in 3.8.

[jakkdl]

Yay!

Replacing the open_memory_channel definition with the class-based one won't work until 3.9+, but I realized I could get the best of both worlds with an if TYPE_CHECKING guard. (I first thought I could write a .pyi only for that function, but that'd override the whole file).

I also wrote some tests using https://github.com/davidfritzsche/pytest-mypy-testing to check for expected mypy error outputs, but it's kinda broken and unmaintained, and I failed to get it working with skipping on earlier versions, so not gonna try and get that working in the CI. But I did get it passing locally at least. It might be worth trying to get trio-typings testing infrastructure replicated here.

[jakkdl]

also rebased and squashed, but the return type of wait_task_rescheduled still needs to be figured out.

[A5rocks]

According to the docs, wait_task_rescheduled returns whatever trio.lowlevel.reschedule is passed. I think Any (or object but pragmatism > type safety) would be best.

Maybe we should brainstorm a nicer interface for this? (we could call it a one-shot channel /s :^) But we should probably just make an issue instead of making one on the spot.

[A5rocks]

A couple last comments, I think everything else is done.

[jakkdl]

According to the docs, wait_task_rescheduled returns whatever trio.lowlevel.reschedule is passed. I think Any (or object but pragmatism > type safety) would be best.

Maybe we should brainstorm a nicer interface for this? (we could call it a one-shot channel /s :^) But we should probably just make an issue instead of making one on the spot.

This is definitely above my paygrade in terms of experience or opinions on trio, so a separate issue to evaluate it sounds good to me if you want to tackle it.

[A5rocks]

Looks good, someone else should take a look and merge if it looks good to them too.

[jakkdl]

A big thank you to @A5rocks for the very helpful and detailed reviews! 🚀