Annotate the test_choices module

This required a bit of refactoring to get the type of `STATUS` correct
for each suite.

There are two cases which I decided not to support in the type system:
- passing a list instead of a tuple when defining an option group
- `in` checks using a data type that doesn't match the choices

tests/test_choices.py

@@ -1,77 +1,88 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+import pytest
 from django.test import TestCase
 
 from model_utils import Choices
 
+T = TypeVar("T")
 
-class ChoicesTests(TestCase):
- def setUp(self) -> None:
- self.STATUS = Choices('DRAFT', 'PUBLISHED')
 
- def test_getattr(self) -> None:
- self.assertEqual(self.STATUS.DRAFT, 'DRAFT')
+class ChoicesTestsMixin(Generic[T]):
 
- def test_indexing(self) -> None:
- self.assertEqual(self.STATUS['PUBLISHED'], 'PUBLISHED')
+ STATUS: Choices[T]
 
- def test_iteration(self) -> None:
- self.assertEqual(tuple(self.STATUS),
- (('DRAFT', 'DRAFT'), ('PUBLISHED', 'PUBLISHED')))
-
- def test_reversed(self) -> None:
- self.assertEqual(tuple(reversed(self.STATUS)),
- (('PUBLISHED', 'PUBLISHED'), ('DRAFT', 'DRAFT')))
+ def test_getattr(self) -> None:
+ assert self.STATUS.DRAFT == 'DRAFT'
 
  def test_len(self) -> None:
- self.assertEqual(len(self.STATUS), 2)
+ assert len(self.STATUS) == 2
 
  def test_repr(self) -> None:
- self.assertEqual(repr(self.STATUS), "Choices" + repr((
+ assert repr(self.STATUS) == "Choices" + repr((
  ('DRAFT', 'DRAFT', 'DRAFT'),
  ('PUBLISHED', 'PUBLISHED', 'PUBLISHED'),
- )))
+ ))
 
  def test_wrong_length_tuple(self) -> None:
- with self.assertRaises(ValueError):
- Choices(('a',))
-
- def test_contains_value(self) -> None:
- self.assertTrue('PUBLISHED' in self.STATUS)
- self.assertTrue('DRAFT' in self.STATUS)
-
- def test_doesnt_contain_value(self) -> None:
- self.assertFalse('UNPUBLISHED' in self.STATUS)
+ with pytest.raises(ValueError):
+ Choices(('a',)) # type: ignore[arg-type]
 
  def test_deepcopy(self) -> None:
  import copy
- self.assertEqual(list(self.STATUS),
- list(copy.deepcopy(self.STATUS)))
+ assert list(self.STATUS) == list(copy.deepcopy(self.STATUS))
 
  def test_equality(self) -> None:
- self.assertEqual(self.STATUS, Choices('DRAFT', 'PUBLISHED'))
+ assert self.STATUS == Choices('DRAFT', 'PUBLISHED')
 
  def test_inequality(self) -> None:
- self.assertNotEqual(self.STATUS, ['DRAFT', 'PUBLISHED'])
- self.assertNotEqual(self.STATUS, Choices('DRAFT'))
+ assert self.STATUS != ['DRAFT', 'PUBLISHED']
+ assert self.STATUS != Choices('DRAFT')
 
  def test_composability(self) -> None:
- self.assertEqual(Choices('DRAFT') + Choices('PUBLISHED'), self.STATUS)
- self.assertEqual(Choices('DRAFT') + ('PUBLISHED',), self.STATUS)
- self.assertEqual(('DRAFT',) + Choices('PUBLISHED'), self.STATUS)
+ assert Choices('DRAFT') + Choices('PUBLISHED') == self.STATUS
+ assert Choices('DRAFT') + ('PUBLISHED',) == self.STATUS
+ assert ('DRAFT',) + Choices('PUBLISHED') == self.STATUS
 
  def test_option_groups(self) -> None:
- c = Choices(('group a', ['one', 'two']), ['group b', ('three',)])
- self.assertEqual(
- list(c),
- [
- ('group a', [('one', 'one'), ('two', 'two')]),
- ('group b', [('three', 'three')]),
- ],
- )
+ # Note: The implementation accepts any kind of sequence, but the type system can only
+ # track per-index types for tuples.
+ if TYPE_CHECKING:
+ c = Choices(('group a', ['one', 'two']), ('group b', ('three',)))
+ else:
+ c = Choices(('group a', ['one', 'two']), ['group b', ('three',)])
+ assert list(c) == [
+ ('group a', [('one', 'one'), ('two', 'two')]),
+ ('group b', [('three', 'three')]),
+ ]
+
+
+class ChoicesTests(TestCase, ChoicesTestsMixin[str]):
+ def setUp(self) -> None:
+ self.STATUS = Choices('DRAFT', 'PUBLISHED')
+
+ def test_indexing(self) -> None:
+ self.assertEqual(self.STATUS['PUBLISHED'], 'PUBLISHED')
+
+ def test_iteration(self) -> None:
+ self.assertEqual(tuple(self.STATUS),
+ (('DRAFT', 'DRAFT'), ('PUBLISHED', 'PUBLISHED')))
+
+ def test_reversed(self) -> None:
+ self.assertEqual(tuple(reversed(self.STATUS)),
+ (('PUBLISHED', 'PUBLISHED'), ('DRAFT', 'DRAFT')))
+
+ def test_contains_value(self) -> None:
+ self.assertTrue('PUBLISHED' in self.STATUS)
+ self.assertTrue('DRAFT' in self.STATUS)
+
+ def test_doesnt_contain_value(self) -> None:
+ self.assertFalse('UNPUBLISHED' in self.STATUS)
 
 
-class LabelChoicesTests(ChoicesTests):
+class LabelChoicesTests(TestCase, ChoicesTestsMixin[str]):
  def setUp(self) -> None:
  self.STATUS = Choices(
  ('DRAFT', 'is draft'),
@@ -157,10 +168,16 @@ def test_composability(self) -> None:
  )
 
  def test_option_groups(self) -> None:
- c = Choices(
- ('group a', [(1, 'one'), (2, 'two')]),
- ['group b', ((3, 'three'),)]
- )
+ if TYPE_CHECKING:
+ c = Choices[int](
+ ('group a', [(1, 'one'), (2, 'two')]),
+ ('group b', ((3, 'three'),))
+ )
+ else:
+ c = Choices(
+ ('group a', [(1, 'one'), (2, 'two')]),
+ ['group b', ((3, 'three'),)]
+ )
  self.assertEqual(
  list(c),
  [
@@ -170,7 +187,7 @@ def test_option_groups(self) -> None:
  )
 
 
-class IdentifierChoicesTests(ChoicesTests):
+class IdentifierChoicesTests(TestCase, ChoicesTestsMixin[int]):
  def setUp(self) -> None:
  self.STATUS = Choices(
  (0, 'DRAFT', 'is draft'),
@@ -216,10 +233,10 @@ def test_doesnt_contain_value(self) -> None:
  self.assertFalse(3 in self.STATUS)
 
  def test_doesnt_contain_display_value(self) -> None:
- self.assertFalse('is draft' in self.STATUS)
+ self.assertFalse('is draft' in self.STATUS) # type: ignore[operator]
 
  def test_doesnt_contain_python_attr(self) -> None:
- self.assertFalse('PUBLISHED' in self.STATUS)
+ self.assertFalse('PUBLISHED' in self.STATUS) # type: ignore[operator]
 
  def test_equality(self) -> None:
  self.assertEqual(self.STATUS, Choices(
@@ -268,10 +285,16 @@ def test_composability(self) -> None:
  )
 
  def test_option_groups(self) -> None:
- c = Choices(
- ('group a', [(1, 'ONE', 'one'), (2, 'TWO', 'two')]),
- ['group b', ((3, 'THREE', 'three'),)]
- )
+ if TYPE_CHECKING:
+ c = Choices[int](
+ ('group a', [(1, 'ONE', 'one'), (2, 'TWO', 'two')]),
+ ('group b', ((3, 'THREE', 'three'),))
+ )
+ else:
+ c = Choices(
+ ('group a', [(1, 'ONE', 'one'), (2, 'TWO', 'two')]),
+ ['group b', ((3, 'THREE', 'three'),)]
+ )
  self.assertEqual(
  list(c),
  [
@@ -284,7 +307,7 @@ def test_option_groups(self) -> None:
 class SubsetChoicesTest(TestCase):
 
  def setUp(self) -> None:
- self.choices = Choices(
+ self.choices = Choices[int](
  (0, 'a', 'A'),
  (1, 'b', 'B'),
  )