Overview

Astrolabe is a plugin for mypy. It checks the column names and the data types of Polars frames and expressions before the code runs. A missing column or a wrong data type is a type error in your editor, not an exception at run time.

What it checks

Declare the columns of a frame once. The checker then knows the schema of the frame at each step of the code, and it reports each step that Polars refuses:

pipeline.pymypy: 3 errors

1import polars as pl

2

3from astrolabe import Columns, TypedDataFrame

4

5

6class Trades(Columns):

7 symbol: pl.String

8 price: pl.Float64

9 quantity: pl.Int64

10

11

12trades: TypedDataFrame[Trades]

13trades.select("venue")

↑error: Unknown column "venue" [unknown-column]

14trades.select(pl.col("symbol") - pl.col("price"))

↑error: "symbol" has no result type at this frame's column types: arithmetic on dtypes String and Float64 is not allowed, and polars raises InvalidOperationError [unsupported-dtype-op]

15trades.filter(pl.col("price"))

↑error: Column "price" is not Boolean, cannot .filter() [wrong-column-shape]

The three errors are three different faults:

• select reads a column that the frame does not have.
• The subtraction mixes a String and a Float64, which Polars does not allow.
• filter needs a Boolean condition, and price is a Float64.

Each operation gives a frame with a new schema. The checker computes that schema, so each later step is checked against it:

pipeline.pymypy: no errors

1class Trades(Columns):

2 symbol: pl.String

3 price: pl.Float64

4 quantity: pl.Int64

5

6

7trades: TypedDataFrame[Trades]

8priced = trades.with_columns(notional=pl.col("price") * pl.col("quantity"))

Resultpriceddeclaredadded

symbolStringpriceFloat64quantityInt64notionalFloat64

Revealed type

TypedDataFrame[TypedDict({'symbol': String, 'price': Float64, 'quantity': Int64, 'notional': Float64}, closed=True), None]

The checker covers the operations that change a schema, for example select, with_columns, join, group_by, pivot and concat. It covers the expressions inside them, at the data types of the frame. A function can declare the columns that it needs, and accept a frame with other columns too. See Frame operations and Open frames.

What it does not check

A type describes the schema of a frame, not its data. The checker does not know:

• the values in a column, for example nulls, ranges or duplicates,
• the number of rows,
• the order of the columns.

For data that comes from outside the program, check the schema at run time with check_schema. A validation library can check the values. See What a type guarantees.

How it works

The plugin runs inside mypy. It reads the code of each expression, and it computes the result of each operation with the typing rules of Polars. mypy keeps the types and reports the errors, so the checks show wherever mypy runs: in your editor, in pre-commit hooks and in CI.

At run time, Astrolabe adds nothing. A typed frame is an ordinary pl.DataFrame, and the code calls Polars as before. See How the checker works.