How to use type hint properly for hypothesis's "stateful testing" example?

Question

I'm trying to properly type hint hypothesis's stateful testing example:

import shutil
import tempfile
from collections import defaultdict

import hypothesis.strategies as st
from hypothesis.database import DirectoryBasedExampleDatabase
from hypothesis.stateful import Bundle, RuleBasedStateMachine, rule


class DatabaseComparison(RuleBasedStateMachine):
    def __init__(self) -> None:
        super().__init__()
        self.tempd = tempfile.mkdtemp()
        self.database = DirectoryBasedExampleDatabase(self.tempd)
        self.model: dict[bytes, set[bytes]] = defaultdict(set)

    keys = Bundle("keys")
    values = Bundle("values")

    @rule(target=keys, k=st.binary())
    def add_key(self, k: bytes) -> bytes:
        return k

    @rule(target=values, v=st.binary())
    def add_value(self, v: bytes) -> bytes:
        return v

    @rule(k=keys, v=values)
    def save(self, k: bytes, v: bytes) -> None:
        self.model[k].add(v)
        self.database.save(k, v)

    @rule(k=keys, v=values)
    def delete(self, k: bytes, v: bytes) -> None:
        self.model[k].discard(v)
        self.database.delete(k, v)

    @rule(k=keys)
    def values_agree(self, k: bytes) -> None:
        assert set(self.database.fetch(k)) == self.model[k]

    def teardown(self) -> None:
        shutil.rmtree(self.tempd)


TestDBComparison = DatabaseComparison.TestCase

However, when I run mypy (Python 3.11.3, mypy 1.3.0):

temp.py:17: error: Need type annotation for "keys"  [var-annotated]
temp.py:18: error: Need type annotation for "values"  [var-annotated]

So how do I fix the "Bundle" thing?

Answer 1

    keys: Bundle[bytes] = Bundle("keys")
    values: Bundle[bytes] = Bundle("values")

should do the trick here!

A Bundle is just a special kind of strategy, which is also usable for @rule(target=<some_bundle>) as well as everywhere a normal strategy would be valid.