Inferred return types

A function can let the checker compute its return type from its body. Write an inferred type as the return type, and wrap each returned value in infer(...). At each call, the result is the union of the types of those values. This is the intended way to type the result of a function that changes an open frame: a remainder marker written by hand repeats what the body does.

Return type	The function returns
InferredExpr	Expressions.
InferredDataFrame	Frames of the class DataFrame.
InferredLazyFrame	Frames of the class LazyFrame.
TypedDataFrame[Schema, InferredRemainder]	A frame with the columns of Schema. The checker computes the remainder.

Return an expression

The type of an expression is long, and a function with several return statements can give several types. InferredExpr gives their union, with no union written by hand:

prices.pymypy: no errors

1import polars as pl

2

3from astrolabe import Columns, InferredExpr, TypedDataFrame, infer

4

5

6class Quotes(Columns):

7 price_usd: pl.Float64

8 price_eur: pl.Float32

9

10

11def price(currency: str) -> InferredExpr:

12 if currency == "USD":

13 return infer(pl.col("price_usd"))

14 return infer(pl.col("price_eur").cast(pl.Float64))

15

16

17quotes: TypedDataFrame[Quotes]

18in_euros = quotes.select(price("EUR"))

Resultin_eurosdeclared

price_usdFloat64

or

price_eurFloat64

Revealed type

TypedDataFrame[TypedDict({'price_usd': Float64}, closed=True), None] | TypedDataFrame[TypedDict({'price_eur': Float64}, closed=True), None]

The checker does not know which branch runs, so the result is either frame.

Return a frame

InferredDataFrame and InferredLazyFrame work in the same way for frames. 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]

Return the declared columns and an inferred remainder

TypedDataFrame[Schema, InferredRemainder] declares the columns of the result, and lets the checker compute the remainder from the body. Below, add_notional accepts any frame with price and quantity, and Rest stands for the other columns of the caller. At the call, the remainder of the result is symbol:

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]

Each returned value must hold each column of the declared schema, with a data type that agrees. The checker reports a value that does not, at its infer(...).

Rules of use

• Write infer(...) around the whole value of a return.
• Use an inferred type only as the return type of a function. Anywhere else, for example on a parameter, it is an error.
• A function with InferredDataFrame must return a DataFrame, and a function with InferredExpr must return expressions.

prices.pymypy: 2 errors

1import polars as pl

2

3from astrolabe import InferredDataFrame, infer

4

5

6def doubled(price: pl.Expr) -> pl.Expr:

7 return infer(price * 2)

↑error: infer() has no effect here: the enclosing function's return type is not InferredExpr/InferredDataFrame/InferredLazyFrame [infer-misuse]

8

9

10def frame() -> InferredDataFrame:

11 return infer(pl.col("price"))

12

13

14frame()

↑error: infer(...) variant is an expression, which is not a valid variant of InferredDataFrame [wrong-inferred-variant]

If one returned value is a plain type that the checker does not track, the whole return type is the plain type. The setting disallow_implicit_inferred_fallback makes that an error. See Configuration.