Skip to content

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

Medium Primitives

Keep the Subclass Type in a Chained Builder with Self

You are writing a query builder whose methods return the builder, so you can chain calls. A subclass adds a page method. With the obvious annotations, mypy rejects a chain that calls page after where:

class QueryBuilder:
    def __init__(self, table: str) -> None:
        self.table = table

    def where(self, clause: str) -> "QueryBuilder":
        return self


class PagedQueryBuilder(QueryBuilder):
    def page(self, number: int, size: int) -> "PagedQueryBuilder":
        return self


PagedQueryBuilder("users").where("active").page(2, 25)
"QueryBuilder" has no attribute "page"

where says it returns a QueryBuilder, so that is the type mypy gives the result, even when you call it on a PagedQueryBuilder. Use Self (from typing) as the return type instead. Self means "the type of the object you call the method on", so the chain keeps its subclass type. This builder only formats a string. It never runs the query.

What to write in QueryBuilder

Each builder has its own table, list of clauses, order column and limit. The clauses start empty, and the order and limit start unset.

  • from_table(cls, table: str) -> Self is a classmethod that returns cls(table), so PagedQueryBuilder.from_table(...) builds a PagedQueryBuilder.
  • where(clause: str) -> Self appends a clause and returns self. Keep the order and any duplicates.
  • order_by(column: str) -> Self sets the order column and returns self. A later call replaces an earlier one.
  • limit(count: int) -> Self sets the limit and returns self. A later call replaces an earlier one, and limit(0) sets a limit of zero.
  • build() -> str returns the query and leaves the builder unchanged:
SELECT * FROM <table> WHERE <c1> AND <c2> ORDER BY <col> LIMIT <n>

Join the parts with one space. Leave out WHERE when there are no clauses, and leave out ORDER BY and LIMIT when they are unset. Copy every string as it is, with no trimming or quoting.

What to write in PagedQueryBuilder

PagedQueryBuilder is a subclass of QueryBuilder.

  • page(number: int, size: int) -> Self stores both numbers, clears any limit set earlier, and returns self. A later call replaces both numbers.
  • Its build has @override, as in the previous exercise. With a page set, it adds LIMIT <size> OFFSET <(number - 1) * size> to the base query. Without one, it returns the base query.
  • Store the numbers as given. Do not check or clamp them.
  • Call limit before page in a chain that uses both. A limit call after page would add a second LIMIT.
q = PagedQueryBuilder.from_table("t")
assert q.where("active").order_by("id").limit(7).page(3, 10) is q
assert q.build() == "SELECT * FROM t WHERE active ORDER BY id LIMIT 10 OFFSET 20"
assert q.build() == q.build()

What solve returns

solve(table, wheres, order, limit, page, per_page) -> tuple[str, str] builds two independent queries for table, one with each class:

  • Add every clause in wheres to both builders.
  • Call order_by(order) on both, only when order is not empty.
  • Call limit(limit) on the plain builder, only when limit > 0.
  • Call page(page, per_page) on the paged builder, only when page > 0.
  • Return (plain.build(), paged.build()).

In solve, assign each chained result back to its variable, as in paged = paged.where(clause). Then paged.page(...) passes mypy only if your methods return Self.

assert solve("t", [], "", 5, 0, 0) == ("SELECT * FROM t LIMIT 5", "SELECT * FROM t")
assert solve("t", [], "", 0, 1, 10) == ("SELECT * FROM t", "SELECT * FROM t LIMIT 10 OFFSET 0")
assert solve("t", ["x=1", "x=1"], "", 0, 0, 0) == (
    "SELECT * FROM t WHERE x=1 AND x=1", "SELECT * FROM t WHERE x=1 AND x=1",
)

Going further

Before Python 3.11 added Self, the same annotation needed a type parameter bound to the class. Use a return type of QueryBuilder, not Self, for a method that always builds a plain QueryBuilder. mypy rejects Self there, because a subclass call would get the wrong type.