static type checking for Polars

Catch Polars bugs before your code runs.

Astrolabe tells mypy, the Python type checker, how to track the contents of your Polars dataframes. A missing column or a wrong data type becomes a type error in your editor, not a runtime failure.

Read the documentationInstall

Python 3.12+ · mypy 1.13+ · Polars 1.43.2

Speed and type safety

Don’t compromise. Keep the speed and expressiveness of dataframes, with the correctness guarantees of statically typed code.

Guardrails and context

Types give people and coding agents up-to-date context and requirements, and the checker validates every refactor across all code paths.

Plain Polars, no overhead

You write ordinary Polars, and a typed frame is an ordinary DataFrame at run time. No new API, no cost to your pipeline.

how it works

1

Declare the schema at load time

Astrolabe reads the schema that you pass to Polars, so the frame has typed columns from the start.

pipeline.pymypy · no errors
import polars as pl
trades = pl.read_csv(
"trades.csv",
schema={
"symbol": pl.String,
"price": pl.Float64,
"quantity": pl.Int64,
},
)
  • A plain read_csv call. Nothing here comes from Astrolabe.
  • The schema gives each column a data type. The checker reads it, so trades now has known columns.
2

Typos are caught

A reference to a column that does not exist is a type error in your editor, not an exception at run time.

pipeline.pymypy · 1 error
# trades as in step 1
trades.select("prcie")
^ error: Unknown column "prcie" [unknown-column]
  • select asks for a column named prcie. trades has price, not prcie, so mypy reports the line.
3

Every step updates the schema

Each operation produces a frame with its own schema, and every later use is checked against it.

pipeline.pymypy · 1 error
# trades as in step 1
symbols = trades.select("symbol", "price")
# ... later in the pipeline
symbols.filter(pl.col("quantity") > 0)
^ error: Unknown column "quantity" [unknown-column]
  • select keeps two columns. The checker knows that symbols has no quantity.
  • Later, a filter reads quantity from symbols. The column is gone, and the error points at this line.
4

Data types are checked too

Each operation needs columns of the right data type, and a column of the wrong type is an error, like a missing one.

pipeline.pymypy · 2 errors
# trades as in step 1
venues = pl.read_csv(
"venues.csv",
schema={"symbol": pl.Int64, "venue": pl.String},
)
trades.filter(pl.col("price"))
^ error: Column "price" is not Boolean, cannot .filter() [wrong-column-shape]
trades.join(venues, on="symbol")
^ error: Join key "symbol" is String on the left and Int64 on the right [join-key-dtype]
  • A second file. Here, symbol is an Int64 code, not a String.
  • filter keeps the rows where a condition is true, so it needs a Boolean column. price is a Float64.
  • A join compares the keys on the two sides. A String key cannot match an Int64 key.
5

New columns are typed

The checker works out the data type of each new column, as Polars does at run time.

pipeline.pymypy · no errors
# trades as in step 1
priced = trades.with_columns(
(pl.col("price") * pl.col("quantity")).alias("notional"),
)
by_symbol = priced.group_by("symbol").agg(
pl.col("notional").sum(),
pl.len(),
)
  • A Float64 times an Int64 gives a Float64. priced has a new column, notional.
  • group_by keeps the key and the aggregates. The sum stays a Float64, and pl.len() counts rows as a UInt32.
6

Write generic functions

Astrolabe's rich grammar of generics lets you work with partially known schemas. Annotate a parameter with only the columns that the function reads, and let Rest stand for any others.

features.pymypy · 1 error
from astrolabe import Columns, TypedDataFrame, other_columns
class Priced(Columns):
price: pl.Float64
quantity: pl.Int64
Rest = other_columns()
def notional(df: TypedDataFrame[Priced, Rest]):
return df.select(
(pl.col("price") * pl.col("quantity")).alias("notional"),
)
quotes = pl.read_csv(
"quotes.csv",
schema={"symbol": pl.String, "price": pl.Float64},
)
notional(quotes)
^ error: Missing required column "quantity" [missing-column]
  • A Columns class names columns and their data types, like a dataclass.
  • Rest stands for all the other columns, the ones that the function does not name.
  • The function takes any frame with price and quantity. Because of Rest, the frame can have other columns too.
  • quotes has symbol and price, but no quantity.
  • The call gives the function a frame without a column that it needs. The error is at the call, not inside the function.
7

Let inference track the rest

The schema of each frame becomes part of the function's signature. Declare the columns that the function adds, and inference works out what the body does to the others. Each caller gets its own columns back.

features.pymypy · no errors
from astrolabe import InferredRemainder, infer
class WithNotional(Priced):
notional: pl.Float64
def add_notional(
df: TypedDataFrame[Priced, Rest],
) -> TypedDataFrame[WithNotional, InferredRemainder]:
return infer(
df.with_columns(
notional=pl.col("price") * pl.col("quantity"),
),
)
result = add_notional(trades)
large = result.filter(pl.col("notional") > 1000)
  • WithNotional takes the columns of Priced and adds notional.
  • The function promises the columns of WithNotional. InferredRemainder stands for the other columns: the checker reads them from the body.
  • infer(...) marks the returned value. The checker makes sure that it holds WithNotional.
  • trades brings symbol as an other column. result keeps it, next to the new notional.
  • result now has a notional column of type Float64, so later code can use it like any other column.

pipeline.pymypy · no errors
import polars as pl
trades = pl.read_csv(
"trades.csv",
schema={
"symbol": pl.String,
"price": pl.Float64,
"quantity": pl.Int64,
},
)

what the checker knows about trades

symbol Stringprice Float64quantity Int64

Step 1 of 7

Want to go further?

Learn the full Astrolabe type grammar and the Polars operations it supports.

Read the documentation →

features

Formally defined and proven

The type grammar is a consistent system with a written specification. Its core theorems are proved in Lean, and its data-type rules are tested against Polars.

Broad Polars coverage

Data frames, lazy frames, expressions and selectors, through select, join,group_by, unpivot and more, each with the schema that Polars gives.

Checks at run time, too

For data that comes from outside, check_schema verifies a frame against its declared type at run time, andforce_schema casts it to that type.

Set up in one command

astrolabe init adds the plugin to your mypy configuration in one step. It then runs wherever mypy runs: your editor, pre-commit hooks and CI.

Honest and pragmatic

Not every schema is known before the code runs, and that’s fine. The checker marks such a frame or column as unknown, and every other check still applies.

Choose your strictness

One setting decides how strict the checker is with steps it cannot fully verify: lenient for an existing codebase, strict for new code, adjustable per operation.

how Astrolabe compares

Validation libraries check data at run time, only where the code calls a check, and only on the paths that run. Astrolabe analyzes every code path and every step, without running your code.

Astrolabepanderapatitodataframely
When mypy runs
Every column referenceA column that the frame does not have is an error wherever the code reads it: in select, pl.col, joins and group keys.✓–––
Schema tracked through every transformEach operation updates the schema, so the code is checked at every step, not only where a frame enters a function.✓–––
Data types in operationsA filter needs a Boolean, join keys must have the same type, and arithmetic must fit the column types.✓–––
Function arguments checked by their columnsA call passes when the frame has the columns that the parameter declares, wherever the frame comes from.✓◐checks the schema label, not the columns–◐checks the schema label, not the columns
Functions that accept extra columnsA function names only the columns it reads. The caller’s other columns pass through to the result, still typed.✓–––
When the code runs
No required overheadType safety adds nothing to the pipeline at run time: no validation pass and no wrapper around Polars.✓–––
Columns and data typesCheck a frame from outside against its declared type.✓check_schema✓✓✓
Cast a frame to its schemaConvert each column to its declared data type.✓force_schema✓✓✓
ValuesNulls, ranges, uniqueness and custom rules on the data.–✓✓✓

Checked with mypy 2.4 against pandera 0.35, patito 0.8 and dataframely 3.1.

Install

Requires Python 3.12 or later.

# add the package
$ uv add git+https://github.com/thomasblanchet/astrolabe
# register the plugin with mypy
$ uv run astrolabe init
# add the package
$ pip install git+https://github.com/thomasblanchet/astrolabe
# register the plugin with mypy
$ astrolabe init
# add the package
$ poetry add git+https://github.com/thomasblanchet/astrolabe
# register the plugin with mypy
$ poetry run astrolabe init