4. Let the body decide the return type

A function that adds a column to an open frame returns the columns of the caller, and the new one. The return type must say both.

Below, add_notional accepts any frame with the columns of Priced: price and quantity. Rest stands for the other columns of the caller. The function adds notional, so its result has the columns of WithNotional and the other columns of the caller.

Let the checker compute the remainder

Write InferredRemainder in the place of the remainder, and wrap the returned value in infer(...). The checker computes the remainder from the body:

features.pymypy: no errors

1import polars as pl

2

3from astrolabe import Columns, InferredRemainder, TypedDataFrame, infer, other_columns

4

5

6class Priced(Columns):

7 price: pl.Float64

8 quantity: pl.Int64

9

10

11class WithNotional(Priced):

12 notional: pl.Float64

13

14

15Rest = other_columns()

16

17

18def add_notional(

19 df: TypedDataFrame[Priced, Rest],

20) -> TypedDataFrame[WithNotional, InferredRemainder]:

21 return infer(df.with_columns(notional=pl.col("price") * pl.col("quantity")))

22

23

24class Trades(Columns):

25 symbol: pl.String

26 price: pl.Float64

27 quantity: pl.Int64

28

29

30trades: TypedDataFrame[Trades]

31result = add_notional(trades)

Resultresultdeclaredaddedfrom the remainder

priceFloat64quantityInt64notionalFloat64symbolString

Revealed type

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

The returned value must hold each column of WithNotional, with the same data type. At each call, the result keeps the other columns of the caller: the frame of Trades has symbol, and so does result. If a caller passes a notional column, with_columns replaces it, and the checker takes that into account.

Let the checker compute the whole frame

InferredDataFrame lets the checker compute the whole return type. With two return statements, the result is the union of the two frames:

features.pymypy: no errors

1class Priced(Columns):

2 price: pl.Float64

3 quantity: pl.Int64

4

5

6class Trades(Columns):

7 symbol: pl.String

8 price: pl.Float64

9 quantity: pl.Int64

10

11

12Rest = other_columns()

13trades: TypedDataFrame[Trades]

14

15

16def summary(df: TypedDataFrame[Priced, Rest], by_value: bool) -> InferredDataFrame:

17 if by_value:

18 return infer(df.select((pl.col("price") * pl.col("quantity")).sum().alias("value")))

19 return infer(df.select(pl.col("quantity").sum()))

20

21

Resultsummary(trades, True)declared

valueFloat64

or

quantityInt64

Revealed type

TypedDataFrame[TypedDict({'value': Float64}, closed=True), None] | TypedDataFrame[TypedDict({'quantity': Int64}, closed=True), None]

Write infer(...) around the whole value of each return. See Inferred return types for InferredExpr and InferredLazyFrame.

Next: 5. Declare reusable expressions.