Overload return type of recv

[BenjyWiener]

Currently recv and recv_streaming are typed as always returning Data, even if decode is set. This means that to satisfy type checkers, we need to use isinstance, cast, etc.

async for chunk in ws.recv_streaming(decode=False):
    assert isinstance(chunk, bytes)
    ...

I propose using typing.overload to narrow the return type when decode is not None:

@overload
async def recv(self, decode: Literal[True]) -> str:
    ...
@overload
async def recv(self, decode: Literal[False]) -> bytes:
    ...
@overload
async def recv(self, decode: None) -> Data:
    ...
async def recv(self, decode: bool | None = None) -> Data:
    """(actual implementation)"""

Alternatively, new recv_str/recv_bytes methods can automatically raise an Exception if the received data is not of the expected type.

[aaugustin]

Yes, this makes sense to me.

[aaugustin]

Worth checking for other instances of the same problem throughout the library.

[aaugustin]

I tried it but mypy doesn't survive the combination of async iteration and type overloads, as you can see in the CI run:

mypy --strict src
src/websockets/asyncio/messages.py:199: error: Overloaded function implementation cannot produce return type of signature 1  [misc]
src/websockets/asyncio/messages.py:199: error: Overloaded function implementation cannot produce return type of signature 2  [misc]
src/websockets/asyncio/messages.py:199: error: Overloaded function implementation cannot produce return type of signature 3  [misc]
src/websockets/asyncio/connection.py:339: error: Overloaded function implementation cannot produce return type of signature 1  [misc]
src/websockets/asyncio/connection.py:339: error: Overloaded function implementation cannot produce return type of signature 2  [misc]
src/websockets/asyncio/connection.py:339: error: Overloaded function implementation cannot produce return type of signature 3  [misc]
src/websockets/asyncio/connection.py:384: error: "Coroutine[Any, Any, AsyncIterator[str | bytes]]" has no attribute "__aiter__" (not async iterable)  [attr-defined]
src/websockets/asyncio/connection.py:384: note: Maybe you forgot to use "await"?
Found 7 errors in 2 files (checked 45 source files)

This looks like python/mypy#10301.

[aaugustin]

Using the workaround suggested in that mypy issue appears to work.

[aaugustin]

Would you be able to confirm that #1579 does what you want, please?

Especially if you're using recv_streaming as I'm somewhat puzzled by the mypy workaround...

[BenjyWiener]

From what I understand, this isn't so much a Mypy bug as a weird side-effect of the way yield affects the return type of async functions. (Namely, stopping it from getting wrapped in Coroutine.)

I played around in the Mypy playground, as well as with Pyright, and removing async from the overload definitions seems to work.

Also note that I removed the default values from the overloads, otherwise Pyright complains about overlapping overloads.

[aaugustin]

When the default values aren't included in the overloads — i.e. when I follow your original suggestion — I get this error:

src/websockets/asyncio/connection.py:242: error: All overload variants of "recv" of "Connection" require at least one argument  [call-overload]
src/websockets/asyncio/connection.py:242: note: Possible overload variants:
src/websockets/asyncio/connection.py:242: note:     def recv(self, decode: Literal[True]) -> Coroutine[Any, Any, str]
src/websockets/asyncio/connection.py:242: note:     def recv(self, decode: Literal[False]) -> Coroutine[Any, Any, bytes]
src/websockets/asyncio/connection.py:242: note:     def recv(self, decode: None) -> Coroutine[Any, Any, str | bytes]
src/websockets/asyncio/connection.py:303: error: Argument 1 to "get" of "Assembler" has incompatible type "bool | None"; expected "None"  [arg-type]
src/websockets/asyncio/connection.py:378: error: Argument 1 to "get_iter" of "Assembler" has incompatible type "bool | None"; expected "None"  [arg-type]
src/websockets/sync/connection.py:240: error: All overload variants of "recv" of "Connection" require at least one argument  [call-overload]
src/websockets/sync/connection.py:240: note: Possible overload variants:
src/websockets/sync/connection.py:240: note:     def recv(self, timeout: float | None, decode: Literal[True]) -> str
src/websockets/sync/connection.py:240: note:     def recv(self, timeout: float | None, decode: Literal[False]) -> bytes
src/websockets/sync/connection.py:240: note:     def recv(self, timeout: float | None, decode: None) -> str | bytes
src/websockets/sync/connection.py:302: error: Argument 2 to "get" of "Assembler" has incompatible type "bool | None"; expected "None"  [arg-type]
src/websockets/sync/connection.py:372: error: Argument 1 to "get_iter" of "Assembler" has incompatible type "bool | None"; expected "None"  [arg-type]

[BenjyWiener]

Weird, I'm not seeing that behavior in the playground. What if you only have a default for the None case?

[BenjyWiener]

Never mind, I was missing an annotation so Mypy wasn't checking the function call. I was able to reproduce your error, and adding a default to the None overload does indeed fix it.

[aaugustin]

I'm still running in trouble with this one:

def get(self, timeout: float | None = None, decode: bool | None = None) -> Data:

• If I remove all the default values, mypy is unhappy.
• If I leave the default values ( = None) in the overload, mypy is happy but you tell me that Pyright isn't.
• If I remove them only for the Literal[True] and Literal[False] case, then I get this error:

mypy --strict src
src/websockets/sync/connection.py:245: error: parameter without a default follows parameter with a default  [syntax]

[aaugustin]

I tried running pyright to come up with overloads that work well in mypy + pyright.

Unfortunately pyright failed at the most basic step:

src/websockets/__init__.py:75:11 - error: "TYPE_CHECKING" is not a known attribute of module "typing"

I spent some time exploring but I don't understand what's failing there.

[BenjyWiener]

I'm not sure about the TYPE_CHECKING error, but I assume Pyright is assuming the wrong root or not treating websockets as a module. It's appears to be looking at websockets/typing.py.

With regards to the parameter without a default follows parameter with a default error, I think we may need something like the following:

@overload
def get(self, timeout: float | None = None, *, decode: Literal[True]) -> str: ...

@overload
def get(self, timeout: float | None, decode: Literal[True]) -> str: ...

@overload
def get(self, timeout: float | None = None, *, decode: Literal[False]) -> bytes: ...

@overload
def get(self, timeout: float | None, decode: Literal[False]) -> bytes: ...

@overload
def get(self, timeout: float | None = None, decode: None = None) -> Data: ...

def get(self, timeout: float | None = None, decode: bool | None = None) -> Data:
    """implementation"""

It's not pretty, but it's probably a good idea to deprecate passing decode as a non-keyword arg anyway.

[aaugustin]

Good catch 🤦

I'll take more time to setup pyright properly and play with it.

[aaugustin]

I still don't understand how to setup pyright. It doesn't appear to behave like it's documented to behave. I filed microsoft/pyright#9752.

[BenjyWiener]

Okay.

Anyway I don't think it necessarily makes sense to get the library to pass both Mypy and Pyright, since they have very different configuration systems. I would just check that downstream projects using websockets type-check properly with both.

[aaugustin]

I agree about the goal. I was just thinking that "type checking properly within websockets" was a good step towards "type checking properly in downstream projects".

[aaugustin]

Getting websockets to pass pyright appears to a be a significant endeavor. I'm not taking this route.

The other problem is python/mypy#7333

Indeed, the only countermeasure for this problem is using a keyword-only argument. Looking into this now.

[aaugustin]

While defensible, making decode a keyword-only argument is backwards incompatible, meaning that it requires a major version bump + release notes.

EDIT: we can do it only in the overload as shown in a comment above 👍