Skip to content

← Precise Type Hints with Literals and Generics step 11 of 26

Medium Primitives

Partition and Chunk a List with Generic Helpers

Write two helpers that work for any element type. partition splits items into the ones that pass a test and the rest. chunk cuts a sequence into lists of a fixed size. Each helper takes a type parameter T, like first[T] in the article "Generic Functions from Zero". So a caller who passes strings gets strings back.

partition(iter([3, 1, 4, 3]), lambda x: x >= 3)  # ([3, 4, 3], [1])
chunk(("a", "b", "c"), 2)                         # [["a", "b"], ["c"]]
chunk([], 3)                                      # []

What to write

Use these signatures, with the types from collections.abc:

from collections.abc import Callable, Iterable, Sequence

def partition[T](
    items: Iterable[T],
    predicate: Callable[[T], bool],
) -> tuple[list[T], list[T]]:
    ...

def chunk[T](items: Sequence[T], size: int) -> list[list[T]]:
    ...

partition:

  • Loop over items once, and call predicate once per item.
  • Append the item to the first list when predicate returns True, and to the second list otherwise.
  • Keep the items in the order they arrive, duplicates included. Do not sort.
  • Do not catch exceptions from the predicate or the iterator.

chunk:

  • Return consecutive lists of size items. The last list may be shorter.
  • Convert each slice with list(...). A slice of a tuple is a tuple, and the result must be a list of lists.
  • Empty input gives [], with no empty list at the end.
  • If size < 1, raise ValueError("size must be >= 1"), even when the input is empty.
  • Leave the input unchanged.

Write the annotations exactly as shown, with T in each. A bare list or Callable fails mypy --strict.

What solve returns

solve(items, threshold, size) returns a tuple of four values:

  1. The items greater than or equal to threshold, from partition.
  2. The other items, from the same call.
  3. chunk(items, size).
  4. The message from a second call, chunk(items, 0), or "" if that call does not raise.

Pass iter(items) to partition, not the list itself. An iterator has no length and no indexing, so passing one shows that partition needs only one pass. Call chunk(items, size) before the chunk(items, 0) check, and outside its try. If size is invalid, that first call's ValueError propagates.

solve([1, 4, 2, 4], 4, 3)
# ([4, 4], [1, 2], [[1, 4, 2], [4]], "size must be >= 1")

Why these input types

partition reads each item once, so Iterable[T] is all it needs, and it accepts generators too. chunk needs len() and slicing, which Sequence[T] provides. Both helpers build new lists, so callers get results they can change while their inputs stay as they were.

Each [T] belongs to its own function, so both helpers can use the name T. mypy fills in T at each call. Chunking strings gives list[list[str]], and partitioning integers gives two list[int] values. The tests use integers only, so try a string example of your own.