How can I provide a type hint for an object returned from the recv_json method?

[wigging]

I'm using pyright to check some code where I receive a dictionary from socket.recv_json(). I get the correct dictionary object but pyright complains that I can't assign dict to the json_data variable such as json_data: dict. The documentation for recv_json has type hints for the return type of which dict is one of the return types. So I don't understand why pyright is complaining about the type hint. Is there another way I should define the type hint?

# client.py

import zmq

def main():
    """Run client."""
    context = zmq.Context()

    # Socket to talk to server
    print("Connecting to server...")
    socket = context.socket(zmq.REQ)
    socket.connect("tcp://localhost:5555")

    # Send request
    socket.send(b"")

    # Get dictionary
    json_data: dict = socket.recv_json()    # <-- pyright error
    print(type(json_data))
    print(json_data)

if __name__ == "__main__":
    main()

Here is the pyright error:

19:23  error   Type "list[Unknown] | str | int | float | dict[Unknown, Unknown]" is not assignable to declared type "dict[Unknown, Unknown]" ​Pyright
                    Type "list[Unknown] | str | int | float | dict[Unknown, Unknown]" is not assignable to type "dict[Unknown, Unknown]"
                      "float" is not assignable to "dict[Unknown, Unknown]"

Also, here is the server code that I'm using with the client code shown above:

# server.py

import zmq

def main():
    """Run server."""
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.bind("tcp://*:5555")

    while True:
        # Wait for request
        socket.recv()

        # Send a dictionary
        d = {"one": 1, "two": 2, "names": ["bart", "homer"]}
        socket.send_json(d)

if __name__ == "__main__":
    main()

[minrk]

I'm no expert on best practices in typing, but I believe this is what typing.cast is for, because code-wise, there's no way for the type checker to know which of the several types recv_json is going to return (other than a JSONable object), but you can cast to what you know it's going to be in practice (the function doesn't do anything, so it is minimal overhead):

json_data: dict = cast(dict, socket.recv_json())

which is equivalent to the (negligibly) more efficient:

json_data: dict = socket.recv_json() # type: ignore

Or you can use an isinstance check to do "type narrowing", and actually be more confident that your code is correct even when the messages you receive aren't what you expect them to be:

json_data = socket.recv_json()
reveal_type(json_data) # dict or list or str, etc.
if not isinstance(json_data, dict):
    raise ValueError(f"Invalid message: {json_data}")
reveal_type(json_data) # dict

You might also be interested in TypedDict for your message schema, so you'll get the typed values of your fields you get out of messages, too, without having to cast again.

[wigging]

Since I know the types associated with the returned dictionary, using cast like shown below amends the pyright error. However, this doesn't read well; it looks messy.

from typing import cast

json_data = cast(dict[str, int | list[str]], socket.recv_json())

Using a TypedDict still causes the pyright error:

from typing import TypedDict

class JsonDict(TypedDict):
    one: int
    two: int
    names: list[str]

json_data: JsonDict = socket.recv_json()

Checking the instance feels like the better approach but the key and value types are still missing from the type hints. For example, pyright shows the json_data as dict[Unknown, Unknown] but at least it doesn't produce type check errors.

json_data = socket.recv_json()

if not isinstance(json_data, dict):
    raise ValueError(f"Invalid message: {json_data}")

You said this regarding the TypedDict approach:

You might also be interested in TypedDict for your message schema, so you'll get the typed values of your fields you get out of messages, too, without having to cast again.

What do you mean by "message schema"? How would I use it without having to cast again?

[minrk]

Using a TypedDict still causes the pyright error:

You still always need to cast or type:ignore to convert from one type to another, so:

json_data = cast(JSONDict, socket.recv_json())
# or
json_data: JSONDict = socket.recv_json() # type: ignore

but the key and value types are still missing from the type hints.

Yes, this is precisely what TypedDict is for, to annotate a dict and all its keys and values with a single annotation.

What do you mean by "message schema"? How would I use it without having to cast again?

Sorry, I think you understood the concept even if my words weren't clear. I only meant what you did in your example, which is to specify which fields are defined and what types they have in a TypedDict declaration.

So in:

class Message(TypedDict):
    one: int
    two: int
    names: list[str]

json_data = cast(Message, socket.recv_json())
reveal_type(json_data['one']) # int
print(json_data['three']) # error, no such key 'three'

However, this doesn't read well; it looks messy.

Yeah, typed code generally looks messier than untyped code. You kind of have to accept that if you want to use types in Python (I don't like typed Python very much, and don't use it in most of my projects).

But if your objection is because of putting verbose type declarations in-line, remember that type annotations are regular Python objects, and you can store them in variables, consolidate them in utility modules, etc. like anything else in Python:

_my_type = dict[str, int | list[str]]
...
result = cast(_my_type, something())

[wigging]

Pyright complains about this too. Apparently it doesn't allow a variable to be used in a type expression.

json_type = dict[str, int | list[str]]
json_data = cast(json_type, socket.recv_json())    # <-- pyright error

[wigging]

I'm going to use the TypedDict approach (see below) since it defines the types associated with each item in the dictionary.

from typing import TypedDict, cast
import zmq

class JsonDict(TypedDict):
    one: int
    two: int
    names: list[str]

def main():
    """Run client."""
    context = zmq.Context()

    # Socket to talk to server
    socket = context.socket(zmq.REQ)
    socket.connect("tcp://localhost:5555")

    # Send request
    socket.send(b"")

    # Receive dictionary
    json_data = cast(JsonDict, socket.recv_json())
    one_value = json_data["one"]
    names_values = json_data["names"]

    # Print results
    print("Received json_data:", json_data)
    print("Received one item:", one_value)
    print("Received names list:", names_values)

if __name__ == "__main__":
    main()

[minrk]

Great! Glad it works.

Apparently it doesn't allow a variable to be used in a type expression.

It does, but for some reason variables storing types must be defined at the module level (not within the function):

json_type = dict[str, int | list[str]]

def foo() -> None:
    json_data = cast(json_type, socket.recv_json())
    reveal_type(json_data) # dict[str, ...]

[wigging]

Well that's good to know. Thanks again for the help.