Using thiserror

use std::io;
use std::num::ParseIntError;
use thiserror::Error;

#[derive(Debug, Error)]
enum AppError {
    // #[from] provides From<ParseIntError> and establishes the source relationship.
    #[error("invalid port number")]
    NotNumber(#[from] ParseIntError),

    #[error("port {0} is not allowed")]
    InvalidPort(u16),

    // #[source] explicitly defines io::Error as the source.
    #[error("failed to load configuration")]
    Config {
        #[source]
        source: io::Error,
    },

    // #[error(transparent)] delegates representation to io::Error.
    // #[from] provides From<io::Error> and establishes the source relationship.
    #[error(transparent)]
    Io(#[from] io::Error),
}

Roughly equivalent handwritten implementation

use std::error::Error;
use std::fmt;
use std::io;
use std::num::ParseIntError;

#[derive(Debug)]
enum AppError {
    NotNumber(ParseIntError),
    InvalidPort(u16),
    Config { source: io::Error },
    Io(io::Error),
}

// Equivalent to #[error("...")]
impl fmt::Display for AppError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotNumber(_) => {
                write!(f, "invalid port number")
            }

            Self::InvalidPort(port) => {
                write!(f, "port {} is not allowed", port)
            }

            Self::Config { .. } => {
                write!(f, "failed to load configuration")
            }

            // Equivalent to #[error(transparent)]
            Self::Io(err) => err.fmt(f),
        }
    }
}

// Equivalent to the Error implementation generated by:
// #[derive(Error)] together with #[source] and #[from].
impl Error for AppError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            // #[from] implies #[source]
            Self::NotNumber(err) => Some(err),

            Self::InvalidPort(_) => None,

            // Explicit #[source]
            Self::Config { source } => Some(source),

            // #[from] implies #[source]
            // #[error(transparent)] delegates representation to the inner error.
            Self::Io(err) => Some(err),
        }
    }
}

// Generated by #[from] on the ParseIntError variant.
impl From<ParseIntError> for AppError {
    fn from(err: ParseIntError) -> Self {
        Self::NotNumber(err)
    }
}

// Generated by #[from] on the io::Error variant.
impl From<io::Error> for AppError {
    fn from(err: io::Error) -> Self {
        Self::Io(err)
    }
}

Puzzle 1: Trait Default Methods

RangeSetBlaze works with sets of integers such as u8, i16, and other integer types. Each integer type must provide fundamental operations like min_value() and max_value(). At the same time, we want all integer types to share additional behavior, such as an exhausted_range() method.

classDiagram
    direction TB

    class Integer {
        <<abstract class>>
        +min_value() Self // required
        +max_value() Self // required
        +exhausted_range() RangeInclusive~Self~  // code
    }

    class u8 {
        <<concrete class>>
        +min_value() Self
        +max_value() Self
        +exhausted_range() RangeInclusive~Self~  // inherited
    }

    class i16 {
        <<concrete class>>
        +min_value() Self
        +max_value() Self
        +exhausted_range() RangeInclusive~Self~  // inherited
    }

    Integer <|-- u8 : is-a
    Integer <|-- i16 : is-a

In Rust, traits are often thought of as contracts: they define a set of required methods that implementors must provide. In that sense, a trait specifies what functionality a type must have.

However, traits can do more than define interfaces. They can also provide concrete method implementations. When a trait includes such methods, all implementors automatically gain access to them unless they choose to override the default behavior. This gives us something conceptually similar to inheritance: shared behavior defined once and reused across multiple types.

use std::ops::RangeInclusive;

// TECHNIQUE NAME: Trait Default Methods.
trait Integer: Copy + Ord {
    fn min_value() -> Self;
    fn max_value() -> Self;

    // Default behavior inherited by implementors.
    // Any impl can override this method.

    /// Returns an exhausted (empty) range`.
    fn exhausted_range() -> RangeInclusive<Self> {
        debug_assert!(Self::min_value() < Self::max_value(), "Precondition");
        Self::max_value()..=Self::min_value()
    }
}

impl Integer for u8 {
    fn min_value() -> Self {
        u8::MIN
    }

    fn max_value() -> Self {
        u8::MAX
    }
}

impl Integer for i16 {
    fn min_value() -> Self {
        i16::MIN
    }

    fn max_value() -> Self {
        i16::MAX
    }
}

fn main() {
    let r1 = u8::exhausted_range();
    let r2 = i16::exhausted_range();

    assert_eq!(r1, 255..=0);
    assert!(r2.is_empty());
}

// TECHNIQUE NAME (again): Trait Default Methods.

The key idea is that u8 and i16 only need to implement the required methods, min_value() and max_value(). The exhausted_range() method is defined once in the trait and automatically shared by all implementors.

Puzzle 2: Supertrait

A servo is an electric motor that can move to a specified angle. We want a ServoEsp (our type that controls a servo on an ESP32 microcontroller) to work with any code that needs a servo. A ServoPlayerEsp is similar, but with animation ability. (Inspired by device-envoy.)

classDiagram
    direction TB

    class Servo {
        <<abstract class>>
        +set_degrees(degrees)
    }

    class ServoPlayer {
        <<abstract class>>
        +set_degrees(degrees) // from Servo
        +animate(steps)
    }

    class ServoEsp {
        <<concrete class>>
        +set_degrees(degrees)
    }

    class ServoPlayerEsp {
        <<concrete class>>
        +set_degrees(degrees)
        +animate(steps)
    }

    Servo <|-- ServoPlayer : is-a
    Servo <|-- ServoEsp : is-a
    ServoPlayer <|-- ServoPlayerEsp : is-a

In Rust, a trait can depend on another trait. The syntax trait B: A means that B requires A as a prerequisite. In other words, any type that implements B must also implement A.

This is similar to saying that B extends A: anything that works with A will also work with B, but B provides additional functionality on top.

In traditional object-oriented terms, this feels like an “is-a” relationship where B can do everything A can do, plus more, while still being treated as an A when needed.

use std::thread;
use std::time::Duration;

// `Servo` is an abstract class (a trait), not a concrete driver.
trait Servo {
    fn set_degrees(&self, degrees: u16);
}

// `ServoPlayer` is also abstract. It extends `Servo` (supertrait), so any
// `ServoPlayer` can do everything in `Servo` plus animation.
// TECHNIQUE NAME: Supertraits
trait ServoPlayer: Servo {
    // (degrees, milliseconds to hold at that angle)
    fn animate(&self, steps: &[(u16, u64)]);
}

#[derive(Default)]
// Concrete servo driver (similar naming to the real example): `ServoEsp`.
struct ServoEsp;

impl Servo for ServoEsp {
    fn set_degrees(&self, degrees: u16) {
        println!("[ServoEsp] set angle -> {degrees}°");
    }
}

#[derive(Default)]
// Concrete servo player driver that can animate.
struct ServoPlayerEsp;

impl Servo for ServoPlayerEsp {
    fn set_degrees(&self, degrees: u16) {
        println!("[ServoPlayerEsp] set angle -> {degrees}°");
    }
}

impl ServoPlayer for ServoPlayerEsp {
    fn animate(&self, steps: &[(u16, u64)]) {
        for (degrees, ms) in steps {
            self.set_degrees(*degrees);
            println!("[ServoPlayerEsp] hold for {ms}ms");
            thread::sleep(Duration::from_millis(*ms));
        }
    }
}

// Generic program that only needs a `Servo`.
fn center_servo(servo: &impl Servo) {
    servo.set_degrees(90);
}

// Generic program that needs a `ServoPlayer`.
fn run_wave(player: &impl ServoPlayer) {
    player.animate(&[
        (0, 120),
        (45, 100),
        (90, 100),
        (135, 100),
        (180, 120),
        (135, 100),
        (90, 100),
        (45, 100),
        (0, 120),
    ]);
}

fn main() {
    let servo_esp = ServoEsp::default();
    let servo_player_esp = ServoPlayerEsp::default();

    center_servo(&servo_esp);
    // `ServoPlayer` can do everything `Servo` can!
    center_servo(&servo_player_esp);
    // and more.
    run_wave(&servo_player_esp);
}

The key idea is that ServoPlayer builds on top of Servo rather than replacing it. Any type that implements ServoPlayer must also implement Servo, allowing it to be used wherever a Servo is expected while providing additional capabilities. Supertraits therefore offer a trait-based way to express “is-a” relationships and extend behavior incrementally.

Puzzle 3: Extension Traits

We want to add a new method is_odd() to an existing concrete type usize, that type is defined outside our crate (“foreign”).

classDiagram
    direction TB

    class UsizeExtensions {
        <<abstract class>>
        +is_odd() bool
    }

    class usize {
        <<concrete class>>
        +is_odd() bool // inherited
    }

    UsizeExtensions <|-- usize : is-a

For types that we define ourselves, such as structs and enums, we can add methods through an impl block. However, Rust does not allow us to write an inherent impl for a type defined in another crate, including primitive types like usize.

This restriction prevents different crates from attaching conflicting methods to the same type, which would make method resolution ambiguous.

Fortunately, traits provide a workaround. We can define our own trait, implement it for the foreign type, and thereby extend the type with new methods. From the caller’s perspective, these methods behave much like native methods on the type itself.

This pattern is known as extension traits.

// impl usize {
//     fn is_odd(self) -> bool {
//         self & 1 != 0
//     }
// }
// error[E0390]: cannot define inherent `impl` for primitive types

trait UsizeExtensions {
    fn is_odd(self) -> bool;
}

impl UsizeExtensions for usize {
    fn is_odd(self) -> bool {
        self & 1 != 0
    }
}

// TECHNIQUE NAME: extension traits.

fn main() {
    let count: usize = 7;

    assert!(count.is_odd());
    assert!(!12.is_odd());
}

The key idea is that usize “inherits” the is_odd() method from UsizeExtensions. Although we cannot add methods directly to the foreign type, implementing an extension trait gives us a very similar result from the caller’s perspective.

Puzzle 4: Derive-Generated Implementations

We want a small LedLevel enum type with two values (On, Off) that automatically participates in common behaviors (default value, debugging output, equality/order comparisons, hashing, copy/clone).

classDiagram
    direction TB

    class LedLevel {
        <<concrete class>>
        On
        Off
    }

    class Defaultable {
        <<abstract class>>
        +default()
    }

    class Debuggable {
        <<abstract class>>
        +debug_string()
    }

    class EquatableOrdered {
        <<abstract class>>
        +equals(other)
        +compare(other)
    }

    class Hashable {
        <<abstract class>>
        +hash()
    }

    class CopyableCloneable {
        <<abstract class>>
        +copy()
        +clone()
    }

    Defaultable <|-- LedLevel : is-a
    Debuggable <|-- LedLevel : is-a
    EquatableOrdered <|-- LedLevel : is-a
    Hashable <|-- LedLevel : is-a
    CopyableCloneable <|-- LedLevel : is-a

In Rust, the derive macro is a powerful convenience feature. It allows us to automatically generate implementations for many common traits, saving us from writing repetitive boilerplate code. These derived traits often provide essential functionality such as formatting, comparison, hashing, and cloning.

Beyond the standard library, many third-party crates also extend this idea. Libraries like serde, for example, provide derive support for serialization and deserialization traits, making it easy to integrate types with external systems.

// TECHNIQUE NAME: derive-generated implementation

#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Default)]
enum LedLevel {
    On,
    #[default]
    Off,
}

fn main() {
    let default_level = LedLevel::default();
    let on = LedLevel::On;
    let off = LedLevel::Off;

    assert_eq!(default_level, LedLevel::Off);
    assert_ne!(on, off);
    assert!(off > on);

    // `Copy` + `Clone` come from derive too.
    let copied = on;
    let cloned = off.clone();
    assert_eq!(copied, on);
    assert_eq!(cloned, off);
}

The key idea is that a tiny enum like LedLevel can immediately behave like a fully featured type—supporting comparison, copying, hashing, and defaults—simply by attaching #[derive(...)], without writing any manual implementation code.

Puzzle 5: Deref Method Lookup

We want an HtmlBuffer to feel like a string for everyday method calls, while still being its own distinct type. Note that String is a concrete type with storage, not just an interface/trait/abstract class.

classDiagram
    direction TB

    class String {
        <<concrete class>>
        -bytes
        +push_str()
        +len()
        +as_bytes()
    }

    class HtmlBuffer {
        <<concrete class>>
        +new()
        +push_str() // inherited
        +len() // inherited
        +as_bytes() // inherited
    }

    String <|-- HtmlBuffer : is-a

The Deref trait allows a type to specify what it should dereference into. This enables a powerful feature in Rust: method lookup through deref coercion.

When a method is not found on the original type, Rust will automatically try to dereference the value and search again on the target type. This process can repeat multiple times, meaning deref behavior can be chained across multiple wrapper types.

For example, a value like Box<String> does not implement contains() directly. However, Rust will first dereference the Box<String> into String, and then continue method lookup on String. Since String itself derefs into str, the call can continue one more step until contains() is found. As a result, a method call like the following works even though contains is defined on str:

fn main() {
    let s = Box::new(String::from("hello world"));
    assert!(s.contains("world"));
}

// struct HtmlBuffer;
// impl String for HtmlBuffer {
// }
// Error because String is not a trait.

use std::ops::{Deref, DerefMut};

// HtmlBuffer 'wraps' a String
struct HtmlBuffer(String);

impl HtmlBuffer {
    fn new() -> Self {
        Self(String::new())
    }
}

impl Deref for HtmlBuffer {
    type Target = String;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for HtmlBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

// TECHNIQUE NAME: deref lookup.

fn main() {
    let mut page = HtmlBuffer::new();

    // `push_str` and `len` are String methods found via deref lookup.
    page.push_str("<h1>Hello</h1>");
    page.push_str("<p>Rust</p>");

    assert_eq!(page.len(), 25);
    // Borrow the thing `page` points to and compare.
    assert_eq!(&*page, "<h1>Hello</h1><p>Rust</p>");
}

// Good for `Box`, `Rc`, etc., but not really good here.
// Also see AsRef, From, and Borrow traits.

The key idea is that HtmlBuffer feels like a String even though it is a wrapper type. Through Deref and DerefMut, method calls like len() and push_str() are automatically forwarded to the inner String. This makes the wrapper effectively transparent in everyday use, while still preserving its own distinct identity as a separate type.

Puzzle 6: Blanket Implementations

We want to union any number of sets (RangeSetBlaze<T>). In OO terms, any collection that is an iterable of RangeSetBlaze<T> references should inherit this operation.

classDiagram
    direction TB

    class Iterable {
        <<abstract class>>
        +iterator()
    }

    class RangeSetCollection {
        <<abstract class>>
        +iterator() // from Iterable
        +union()
    }

    class VectorOfRangeSetRefs {
        <<concrete class>>
        +iterator()
        +union() // inherited
    }

    class ArrayOfRangeSetRefs {
        <<concrete class>>
        +iterator()
        +union() // inherited
    }

    class AnyOtherRangeSetCollection {
        <<concrete class>>
        +iterator()
        +union() // inherited
    }

    Iterable <|-- RangeSetCollection : is-a
    RangeSetCollection <|-- VectorOfRangeSetRefs : is-a
    RangeSetCollection <|-- ArrayOfRangeSetRefs : is-a
    RangeSetCollection <|-- AnyOtherRangeSetCollection : is-a

A blanket implementation allows a trait to be implemented for an entire category of types rather than for a single concrete type. Instead of writing many individual implementations, we can describe a set of requirements and automatically grant the trait to every type that satisfies them.

The requirements are typically expressed through trait bounds, often using a where clause. This lets us precisely control which types receive the implementation while keeping the code concise and reusable.

Like other generic features in Rust, blanket implementations are resolved at compile time. The compiler generates the necessary code for each applicable type, providing flexibility without introducing runtime overhead.

use std::collections::BTreeSet;

// For this example, use u64 as our stand-in integer type.
type Integer = u64;

// Mock up RangeSetBlaze as a BTreeSet wrapper for demo purposes.
#[derive(Debug, Clone, PartialEq, Eq)]
struct RangeSetBlaze {
    values: BTreeSet<Integer>,
}

impl RangeSetBlaze {
    fn new() -> Self {
        Self {
            values: BTreeSet::new(),
        }
    }

    // Construct from a slice of Integer values.
    fn from_slice(values: &[Integer]) -> Self {
        Self {
            values: values.iter().copied().collect(),
        }
    }

    // Return a new set that is the union of two borrowed sets.
    fn union(&self, other: &Self) -> Self {
        Self {
            values: self.values.union(&other.values).copied().collect(),
        }
    }

    fn is_empty(&self) -> bool {
        self.values.is_empty()
    }
}

// A RangeSetCollection must be iterable over borrowed RangeSetBlaze values.
// In return, the trait gives it a union() method over those values.
trait RangeSetCollection<'a>: IntoIterator<Item = &'a RangeSetBlaze> {
    fn union(self) -> RangeSetBlaze
    where
        Self: Sized,
    {
        let mut result = RangeSetBlaze::new();
        for set in self {
            result = RangeSetBlaze::union(&result, set);
        }
        result
    }
}

// TECHNIQUE NAME: blanket implementation
//
// This is broader than a normal extension trait impl.
// Instead of adding union() to one type, we add it to every type
// that can be turned into an iterator over &RangeSetBlaze.
impl<'a, I> RangeSetCollection<'a> for I where I: IntoIterator<Item = &'a RangeSetBlaze> {}

fn main() {
    let a = RangeSetBlaze::from_slice(&[1, 2, 3]);
    let b = RangeSetBlaze::from_slice(&[3, 4, 5]);
    let c = RangeSetBlaze::from_slice(&[5, 7, 9]);

    let expected = RangeSetBlaze::from_slice(&[1, 2, 3, 4, 5, 7, 9]);
    // A vector of sets can be unioned together.
    assert_eq!(vec![&a, &b, &c].union(), expected);
    // An array of sets can be unioned together.
    assert_eq!([&a, &b, &c].union(), expected);
    // A filtered option can be unioned.
    assert_eq!(Some(&a).filter(|set| !set.is_empty()).union(), a);
}

Some(&a).filter(|set| !set.is_empty()).union()

The key idea is that union() is written only once, yet it automatically becomes available on many different collection types. In the example, a vector, an array, and even a filtered Option can all call the same union() method because they satisfy the required iterator constraint.

Puzzle 7: Macro-Generated Implementations

Treat 15 integer-like types uniformly: u8, …, usize, …i128, char, IPv4, IPv6

For example, with a min, max, and ability to add one (unchecked).

classDiagram
    direction TB

    class Integer {
        <<abstract class>>
        +add_one()
        +min_value()
        +max_value()
    }

    class NumericInteger {
        <<abstract class>>
        +add_one()
        +min_value()
        +max_value()
    }

    class IpInteger {
        <<abstract class>>
        +add_one()
        +min_value()
        +max_value()
    }

    class CharInteger {
        <<abstract class>>
        +add_one()
        +min_value()
        +max_value()
    }

    class u8 {
        <<concrete class>>
        +add_one() // inherited
        +min_value() // inherited
        +max_value() // inherited
    }

    class IPv4Type {
        <<concrete class>>
        +add_one() // inherited
        +min_value() // inherited
        +max_value() // inherited
    }

    class IPv6Type {
        <<concrete class>>
        +add_one() // inherited
        +min_value() // inherited
        +max_value() // inherited
    }

    class CharType {
        <<concrete class>>
        +add_one() // inherited
        +min_value() // inherited
        +max_value() // inherited
    }

    Integer <|-- NumericInteger : is-a
    Integer <|-- IpInteger : is-a
    Integer <|-- CharInteger : is-a

    NumericInteger <|-- u8 : is-a
    IpInteger <|-- IPv4Type : is-a
    IpInteger <|-- IPv6Type : is-a
    CharInteger <|-- CharType : is-a

In addition to procedural macros like derive, which take a token stream and transform it into another token stream, declarative macros (macro_rules!) can also generate code — including full method implementations.

The key difference here is that instead of manually writing nearly identical impl blocks for each type, we encode the shared logic once and reuse it across all implementations via macros. This significantly reduces repetition while keeping the specialization for each type explicit where needed.

use std::net::{Ipv4Addr, Ipv6Addr};

trait Integer: Copy + Ord {
    fn add_one(self) -> Self;
    fn min_value() -> Self;
    fn max_value() -> Self;
}

macro_rules! impl_integer_ops_num {
    ($t:ty) => {
        fn add_one(self) -> Self {
            self + 1
        }

        fn min_value() -> Self {
            <$t>::MIN
        }

        fn max_value() -> Self {
            <$t>::MAX
        }
    };
}

macro_rules! impl_integer_ops_ip {
    ($ip_type:ty, $representation_type:ty) => {
        fn add_one(self) -> Self {
            <$ip_type>::from(<$representation_type>::from(self) + 1)
        }

        fn min_value() -> Self {
            <$ip_type>::from(<$representation_type>::MIN)
        }

        fn max_value() -> Self {
            <$ip_type>::from(<$representation_type>::MAX)
        }
    };
}

macro_rules! impl_integer_ops_char {
    () => {
        fn add_one(self) -> Self {
            let mut num = u32::from(self) + 1;
            if num == 0xD800 {
                num = 0xE000;
            }
            char::from_u32(num).expect("next char must be a valid Unicode scalar value")
        }

        fn min_value() -> Self {
            char::MIN
        }

        fn max_value() -> Self {
            char::MAX
        }
    };
}

impl Integer for i8 { impl_integer_ops_num!(i8); }
impl Integer for u8 { impl_integer_ops_num!(u8); }
impl Integer for i16 { impl_integer_ops_num!(i16); }
impl Integer for u16 { impl_integer_ops_num!(u16); }
impl Integer for i32 { impl_integer_ops_num!(i32); }
impl Integer for u32 { impl_integer_ops_num!(u32); }
impl Integer for i64 { impl_integer_ops_num!(i64); }
impl Integer for u64 { impl_integer_ops_num!(u64); }
impl Integer for i128 { impl_integer_ops_num!(i128); }
impl Integer for u128 { impl_integer_ops_num!(u128); }
impl Integer for isize { impl_integer_ops_num!(isize); }
impl Integer for usize { impl_integer_ops_num!(usize); }
impl Integer for Ipv4Addr { impl_integer_ops_ip!(Ipv4Addr, u32); }
impl Integer for Ipv6Addr { impl_integer_ops_ip!(Ipv6Addr, u128); }
impl Integer for char { impl_integer_ops_char!(); }

// Can't use num_traits::PrimInt as a 'subclass'
// because Rust worries that 'char' etc. might
// be added to it later (coherence).
// HOWEVER: You could define your own PrimInt with 12 impls.

// TECHNIQUE NAME: macro-generated implementation.

fn main() {
    let x: u8 = 254;
    assert_eq!(x.add_one(), 255);
    assert_eq!('a'.add_one(), 'b');

    assert_eq!(char::min_value(), char::MIN);
    assert_eq!(char::max_value(), char::MAX);
    assert_eq!(Ipv4Addr::min_value(), Ipv4Addr::new(0, 0, 0, 0));
    assert_eq!(Ipv4Addr::max_value(), Ipv4Addr::new(255, 255, 255, 255));
}

The key idea is that macros let us scale a single abstraction (Integer) across many unrelated types by generating repetitive impl blocks automatically, while still allowing special-case logic (like char and IP addresses) where needed.

Puzzle 8: Constraint-Gated Methods

We have abstract class OutputArray<N> with two methods. For N = 8, however, we want one extra method because it maps cleanly to a u8 bitmask.

classDiagram
    direction TB

    class OutputArrayN["OutputArray~N~"] {
        <<abstract class>>
        +new()s
        +set_level_at_index(index, level)
    }

    class OutputArray4["OutputArray~4~"] {
        <<concrete class>>
        +new() // inherited
        +set_level_at_index(index, level) // inh.
    }

    class OutputArray8["OutputArray~8~"] {
        <<concrete class>>
        +new() // inherited
        +set_level_at_index(index, level) // inh.
        +set_from_bits(u8)
    }


    OutputArrayN <|-- OutputArray4 : is-a
    OutputArrayN <|-- OutputArray8 : is-a

Rust allows methods to be conditionally attached to a type by placing them in separate impl blocks with additional constraints. These constraints can take many forms, including const generic values, trait bounds, lifetimes, or combinations of them.

As a result, a type’s API can vary depending on the compile-time properties it satisfies. Rather than exposing every method to every instance of a type, Rust allows certain methods to exist only when specific requirements are met.

#[derive(Debug, Clone, Copy)]
struct OutputArray<const N: usize> {
    levels: [bool; N],
}

impl<const N: usize> OutputArray<N> {
    fn new() -> Self {
        Self { levels: [false; N] }
    }

    fn set_level_at_index(&mut self, index: usize, level: bool) {
        self.levels[index] = level;
    }
}

// Hardcoded 8, matching the bit width expected by a u8 mask.
impl OutputArray<8> {
    fn set_from_bits(&mut self, mut bits: u8) {
        for slot in &mut self.levels {
            *slot = (bits & 1) == 1;
            bits >>= 1;
        }
    }
}

// TECHNIQUE NAME: constraint-gated methods.
// - Put methods in a separate impl block.
// - The methods exist only when the impl's constraints are met.
// - Constraints can be const values, trait bounds, lifetimes, or combinations

fn main() {
    let mut any = OutputArray::<4>::new();
    any.set_level_at_index(2, true);

    let mut eight = OutputArray::<8>::new();
    eight.set_from_bits(0b1011_0001);

    assert_eq!(any.levels, [false, false, true, false]);
    assert_eq!(eight.levels, [true, false, false, false, true, true, false, true]);
}

The key idea is that method availability can depend on compile-time constraints. In this example, only OutputArray<8> gains the set_from_bits() method, because only an 8-element output array can be represented directly by a u8 bitmask. The restriction is enforced entirely at compile time, making invalid usages impossible to compile.

Puzzle 9: Method-Level Constraints

We want two levels of FlashBlock: a top level with new and clear, and a constrained level that adds save and load only when T satisfies both serialization and deserialization capabilities.

classDiagram
    direction TB

    class Serialize {
        <<abstract class>>
        +serialize()
    }

    class Deserialize {
        <<abstract class>>
        +deserialize()
    }

    class TConstrained["T: Serialize + Deserialize"] {
        <<abstract class>>
        +serialize() // inherited
        +deserialize() // inherited
    }

    class WifiCredentials {
        <<concrete class>>
        +serialize() // inherited
        +deserialize() // inherited
    }

    class FlashBlock {
        <<concrete class>>
        +new()
        +clear()
    }

    class FlashBlockForTSerializeDeserialize {
        <<concrete class>>
        +new() // inherited
        +clear() // inherited
        +save(value, key)
        +load(key) T
    }

    Serialize <|-- TConstrained : is-a
    Deserialize <|-- TConstrained : is-a
    TConstrained <|-- WifiCredentials : is-a
    FlashBlock <|-- FlashBlockForTSerializeDeserialize : is-a

Unlike constraint-gated methods, which apply bounds to an entire impl block, Rust allows constraints to be placed directly on individual methods via a where clause. This enables a type to remain generally usable while selectively enabling certain operations only when additional requirements are met.

use std::collections::HashMap;

use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
struct WifiCredentials {
    ssid: String,
    password: String,
}

/// Key point: method-level gating.
/// - The type exists for all use-cases.
/// - Only `save/load<T>` require extra capabilities on `T`.
#[derive(Default)]
struct FlashBlock {
    store: HashMap<String, Vec<u8>>,
}

impl FlashBlock {
    fn new() -> Self {
        Self::default()
    }

    fn clear(&mut self) {
        self.store.clear();
    }

    fn save<T>(&mut self, key: &str, value: &T) -> Result<(), postcard::Error>
    where
        T: Serialize + DeserializeOwned,
    {
        let bytes = postcard::to_stdvec(value)?;
        self.store.insert(key.to_string(), bytes);
        Ok(())
    }

    fn load<T>(&self, key: &str) -> Option<T>
    where
        T: Serialize + DeserializeOwned,
    {
        let bytes = self.store.get(key)?;
        postcard::from_bytes(bytes).ok()
    }
}

// TECHNIQUE NAME: method-level gating.

fn main() -> Result<(), postcard::Error> {
    let mut flash = FlashBlock::new();

    let credentials = WifiCredentials {
        ssid: "HomeWiFi".to_string(),
        password: "secret".to_string(),
    };

    // Works: WifiCredentials satisfies Serialize + Deserialize.
    flash.save("wifi", &credentials)?;
    let loaded: Option<WifiCredentials> = flash.load("wifi");
    let loaded = loaded.unwrap();
    assert_eq!(&loaded.ssid, "HomeWiFi");
    assert_eq!(&loaded.password, "secret");

    flash.clear();
    assert!(flash.load::<WifiCredentials>("wifi").is_none());

    // This compiles, but it asks for the wrong type.
// device-envoy stores type-identifying metadata with each value,
// so this kind of mismatch is likely to be caught at load time.
    flash.save("number", &42u8)?;
    let loaded: Option<String> = flash.load("number");
    assert!(loaded.is_none());

    Ok(())
}

The key idea is that constraints can be attached to individual methods instead of the entire type. In this example, basic operations are always available on FlashBlock, while storage-related operations are only enabled when the value type satisfies the required serialization and deserialization bounds. This keeps the core type simple while making advanced functionality conditional on capability.

Conclusion

Here is the conclusion table from Carl Kadie:

Output Column Names

In Polars, if an explicit name isn’t provided, the output column will take the name of the first expression. This default behavior is generally helpful, but it can lead to unexpected results—especially when the first expression doesn’t directly represent a column. For example, if the first expression is pl.lit(), the resulting column name will be "literal".

This isn’t a big issue when dealing with a single column, but it becomes problematic when using expressions like pl.all() or pl.col("*"). In such cases, logic that depends on column names often breaks.

My current workaround is to tweak the internal logic so that the returned expression doesn’t start with something like pl.lit(), especially in functions that aim to support wildcard expressions.

Expression Context vs DataFrame Context

Since Polars expressions are designed to be evaluated later (not immediately), we cannot rely on any runtime properties of the DataFrame when building them. That means things like the shape or number of rows in the DataFrame are not available during expression construction.

However, some built-in Polars functions help bridge this gap. For instance, pl.len() returns an expression that can be used to represent the number of rows in the future evaluation context. This is useful for writing general-purpose logic in Turtle Island.

A major caveat is that Polars expressions can’t be evaluated directly in Python. You can’t write something like if pl.len() or pl.len() + 1 as regular Python code. These expressions must be used inside other Polars expressions or functions—such as pl.int_range(0, pl.len() + 1)—where they’ll be correctly interpreted and evaluated within the Polars execution context.

Building the Time Machine Core

The goal is to create a WigGT object that behaves like a GT object from Great Tables. As we progressively build the table, we maintain two versions:

• The original (non-interactive) table
• An interactive version that includes a marimo widget (in this case, a slider)

This setup allows us to move back and forth through the table-building steps and see how each method transforms the output — effectively creating a table “time machine.”

Let’s break it down.

def get_allowed_member_names() -> list[str]:
    return ["fmt", "fmt_number", ...]

def lazify(cls: GT) -> GT:
    def add_to_pipeline(func: Callable[..., GT]) -> callable:
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            def inner(self):
                return func(self, *args, **kwargs)

            # container for storing callable objects
            self._pipeline.append(inner)
            return self

        wrapper.__signature__ = inspect.signature(func)
        return wrapper

    for member_name in get_allowed_member_names():
        setattr(
            cls, member_name, add_to_pipeline(getattr(GT, member_name))
        )
    return cls

@lazify
class WigGT:
    def __init__(self, *args, widget, **kwargs):
        self._args = args
        self._widget = widget
        self._kwargs = kwargs
        self.set_to_init()

    @property
    def tables(self) -> list[GT]:
        return list(self._tables)  # return a new list

    def _widgetify(self, obj: GT) -> GT:
        return obj.tab_source_note(html(self._widget))

    def collect(self) -> Self:
        if not self._is_collect:
            new_obj = self._tables[0]  # don't use self._wtables[0]
            for f in self._pipeline:
                new_obj = f(self=copy.copy(new_obj))
                self._tables.append(new_obj)
                self._wtables.append(self._widgetify(new_obj))
            self._is_collect = True
        return self

    def _repr_html_(self) -> str:
        try:
            obj = self._wtables[self._widget.value]
        except IndexError:
            obj = self._wtables[-1]

        if hasattr(obj, "_display_"):
            render_method = "_display_"
        elif hasattr(obj, "_repr_html_"):
            render_method = "_repr_html_"
        elif hasattr(obj, "_mime_"):
            render_method = "_mime_"
        else:
            raise AttributeError(
                "The object does not have a valid render method."
            )
        return getattr(obj, render_method)()  # remember to invoke

    def set_to_init(self) -> None:
        if not getattr(self, "_pipeline", None):
            self._pipeline: list[callable] = []
        else:
            self._pipeline.clear()

        if not getattr(self, "_tables", None):
            self._tables = []
        else:
            self._tables.clear()

        if not getattr(self, "_wtables", None):
            self._wtables = []
        else:
            self._wtables.clear()

        obj = self._make_gt()
        self._tables.append(obj)
        self._wtables.append(self._widgetify(obj))
        self._is_collect = False

    def _make_gt(self) -> GT:
        return GT(*self._args, **self._kwargs)

marimo

Now we’re ready to time travel through the table-building process!

time_widget = mo.ui.slider(start=0, stop=6, step=1, value=0, label="Step")

lazy_wig_gt = (
    WigGT(airquality.head(10).assign(Year=1973), widget=time_widget)
    .opt_stylize(color="pink", style=2)
    .tab_header(
        title="New York Air Quality Measurements",
        subtitle="Daily measurements in New York City (May 1-10, 1973)",
    )
    .tab_spanner(label="Time", columns=["Year", "Month", "Day"])
    .tab_spanner(
        label="Measurement", columns=["Ozone", "Solar_R", "Wind", "Temp"]
    )
    .cols_move_to_start(columns=["Year", "Month", "Day"])
    .cols_label(
        Ozone=html("Ozone,<br>ppbV"),
        Solar_R=html("Solar R.,<br>cal/m<sup>2</sup>"),
        Wind=html("Wind,<br>mph"),
        Temp=html("Temp,<br>&deg;F"),
    )
)
lazy_wig_gt

wig_gt = lazy_wig_gt.collect()
wig_gt

wig_gt.tables

import marimo

__generated_with = "0.14.8"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell
def _():
    import copy
    import inspect
    from functools import wraps
    from typing import Callable, Self

    import pandas as pd
    from great_tables import GT, html
    from great_tables.data import airquality

    return Callable, GT, Self, airquality, copy, html, inspect, wraps


@app.cell
def _(Callable, GT, Self, copy, html, inspect, wraps):
    def get_allowed_member_names() -> list[str]:
        """
        Manually constructing the list —
        it would be great if Great Tables exposed the available method names.
        """
        return [
            "fmt",
            "fmt_number",
            "fmt_integer",
            "fmt_percent",
            "fmt_scientific",
            "fmt_currency",
            "fmt_bytes",
            "fmt_roman",
            "fmt_date",
            "fmt_time",
            "fmt_datetime",
            "fmt_markdown",
            "fmt_image",
            "fmt_icon",
            "fmt_flag",
            "fmt_units",
            "fmt_nanoplot",
            "data_color",
            "sub_missing",
            "sub_zero",
            "opt_stylize",
            "opt_align_table_header",
            "opt_all_caps",
            "opt_footnote_marks",
            "opt_row_striping",
            "opt_vertical_padding",
            "opt_horizontal_padding",
            "opt_table_outline",
            "opt_table_font",
            "cols_align",
            "cols_width",
            "cols_label",
            "cols_move",
            "cols_move_to_start",
            "cols_move_to_end",
            "cols_hide",
            "cols_unhide",
            "tab_header",
            "tab_source_note",
            "tab_spanner",
            "tab_stubhead",
            "tab_style",
            "tab_options",
            "row_group_order",
            "tab_stub",
            "with_id",
            "with_locale",
            "save",
            "show",
            "as_raw_html",
            "write_raw_html",
            "as_latex",
            "pipe",
        ]

    def lazify(cls: GT) -> GT:
        def add_to_pipeline(func: Callable[..., GT]) -> callable:
            @wraps(func)
            def wrapper(self, *args, **kwargs):
                def inner(self):
                    return func(self, *args, **kwargs)

                # container for storing callable objects
                self._pipeline.append(inner)
                return self

            wrapper.__signature__ = inspect.signature(func)
            return wrapper

        for member_name in get_allowed_member_names():
            setattr(
                cls, member_name, add_to_pipeline(getattr(GT, member_name))
            )
        return cls

    @lazify
    class WigGT:
        def __init__(self, *args, widget, **kwargs):
            self._args = args
            self._widget = widget
            self._kwargs = kwargs
            self.set_to_init()

        @property
        def tables(self) -> list[GT]:
            return list(self._tables)  # return a new list

        def _widgetify(self, obj: GT) -> GT:
            return obj.tab_source_note(html(self._widget))

        def collect(self) -> Self:
            if not self._is_collect:
                new_obj = self._tables[0]  # don't use self._wtables[0]
                for f in self._pipeline:
                    new_obj = f(self=copy.copy(new_obj))
                    self._tables.append(new_obj)
                    self._wtables.append(self._widgetify(new_obj))
                self._is_collect = True
            return self

        def _repr_html_(self) -> str:
            try:
                obj = self._wtables[self._widget.value]
            except IndexError:
                obj = self._wtables[-1]

            if hasattr(obj, "_display_"):
                render_method = "_display_"
            elif hasattr(obj, "_repr_html_"):
                render_method = "_repr_html_"
            elif hasattr(obj, "_mime_"):
                render_method = "_mime_"
            else:
                raise AttributeError(
                    "The object does not have a valid render method."
                )
            return getattr(obj, render_method)()  # remember to invoke

        def set_to_init(self) -> None:
            if not getattr(self, "_pipeline", None):
                self._pipeline: list[callable] = []
            else:
                self._pipeline.clear()

            if not getattr(self, "_tables", None):
                self._tables = []
            else:
                self._tables.clear()

            if not getattr(self, "_wtables", None):
                self._wtables = []
            else:
                self._wtables.clear()

            obj = self._make_gt()
            self._tables.append(obj)
            self._wtables.append(self._widgetify(obj))
            self._is_collect = False

        def _make_gt(self) -> GT:
            return GT(*self._args, **self._kwargs)

    return (WigGT,)


@app.cell
def _(mo):
    time_widget = mo.ui.slider(start=0, stop=6, step=1, value=0, label="Step")
    time_widget
    return (time_widget,)


@app.cell
def _(WigGT, airquality, html, time_widget):
    # The lazy_wig_gt object is not interactive until collect() is called.
    lazy_wig_gt = (
        WigGT(airquality.head(10).assign(Year=1973), widget=time_widget)
        .opt_stylize(color="pink", style=2)
        .tab_header(
            title="New York Air Quality Measurements",
            subtitle="Daily measurements in New York City (May 1-10, 1973)",
        )
        .tab_spanner(label="Time", columns=["Year", "Month", "Day"])
        .tab_spanner(
            label="Measurement", columns=["Ozone", "Solar_R", "Wind", "Temp"]
        )
        .cols_move_to_start(columns=["Year", "Month", "Day"])
        .cols_label(
            Ozone=html("Ozone,<br>ppbV"),
            Solar_R=html("Solar R.,<br>cal/m<sup>2</sup>"),
            Wind=html("Wind,<br>mph"),
            Temp=html("Temp,<br>&deg;F"),
        )
    )
    lazy_wig_gt
    return (lazy_wig_gt,)


@app.cell
def _(lazy_wig_gt):
    # The wig_gt object is now interactive
    wig_gt = lazy_wig_gt.collect()
    wig_gt
    return (wig_gt,)


@app.cell
def _(wig_gt):
    # retrieve all `gt` tables
    wig_gt.tables
    return


if __name__ == "__main__":
    app.run()

In Closing

marimo has turned out to be far more powerful than I expected — I’ve had a lot of fun experimenting with it.

I highly recommend checking out their YouTube channel, which is full of high-quality content across many Python topics.

This idea was heavily inspired by their videos — all credit goes to the amazing marimo team!

marimo

1. I asked AI to generate a checklist and wrapped it in a Polars DataFrame called df.
2. I created 10 switch widgets and stacked them into an array widget named status_widgets to represent the status of each checklist item.
3. I extracted the HTML representation of each widget via its _repr_html_() method and inserted it as a new "Status" column in df, which I then wrapped in a Great Tables GT object.
4. I added two source notes using GT.tab_source_note()—one to display progress, and another for a visual progress bar.
5. Finally, I gave the table a nice header with GT.tab_header() and applied some styling using GT.opt_stylize().

Check out the full marimo code below or view it on molab.

import marimo

__generated_with = "0.14.7"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell
def _():
    import polars as pl
    from great_tables import GT, html, md

    return GT, html, pl


@app.cell
def _(pl):
    tasks = [
        "Set DEBUG = False",
        "Configure ALLOWED_HOSTS",
        "Set up a secret key",
        "Collect static files",
        "Apply database migrations",
        "Set up gunicorn or uWSGI",
        "Configure reverse proxy (e.g., Nginx)",
        "Secure the database",
        "Set up HTTPS (SSL)",
        "Configure logging & monitoring",
    ]

    notes = [
        "Never deploy with DEBUG = True ⚠️",
        "Include your domain(s) or IP address 🌐",
        "Use a strong, secure key from an environment variable 🔐",
        "Run `python manage.py collectstatic` 📦",
        "Run `python manage.py migrate` 🗃️",
        "Use as a WSGI server in production 🔄",
        "Serve static/media files and forward to WSGI server 🧭",
        "Use strong credentials, disable remote root login 🛡️",
        "Use Let's Encrypt or your own certificate 🔒",
        "Track errors and app performance 📊",
    ]

    n_row = len(tasks)
    status = ["☐"] * n_row
    data = {"Status": status, "Task": tasks, "Notes": notes}

    df = pl.DataFrame(data)
    return df, n_row


@app.cell
def _(mo, n_row):
    status_widget = mo.ui.switch()
    status_widgets = mo.ui.array([status_widget] * n_row)
    return (status_widgets,)


@app.function
def create_bar(
    x: float,
    max_width: int,
    height: int,
    background_color1: str,
    background_color2: str,
) -> str:
    width = round(max_width * x, 2)
    px_width = f"{width}px"
    return f"""\
    <div style="width: {max_width}px; background-color: {background_color1};">\
        <div style="height:{height}px;width:{px_width};background-color:{background_color2};"></div>\
    </div>\
    """


@app.cell
def _(GT, df, html, n_row, pl, status_widgets):
    done_count = sum(s.value for s in status_widgets)

    gt = (
        GT(
            df.with_columns(
                pl.Series(
                    [status._repr_html_() for status in status_widgets]
                ).alias("Status")
            )
        )
        .tab_source_note(f"{done_count} / {n_row}")
        .tab_source_note(
            html(
                create_bar(
                    done_count / n_row,
                    max_width=750,
                    height=20,
                    background_color1="lightgray",
                    background_color2="#66CDAA",
                )
            )
        )
        .tab_header("✅ Django Deployment Checklist")
        .opt_stylize(color="cyan", style=4)
    )
    gt
    return


if __name__ == "__main__":
    app.run()

marimo

Here, I demonstrate how marimo widgets can be embedded in Great Tables by wrapping them with the html() function provided by the library. This allows interactive widgets to control the table’s appearance in a Quarto environment via WASM — a surprisingly powerful capability, in my opinion.

Check out the full marimo code below or view it on molab.

import marimo

__generated_with = "0.13.15"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    import polars as pl
    from great_tables import GT, html

    return GT, html, mo, pl


@app.cell
def _(pl):
    data = {
        "col1": [2, 5, 7, 10, 15],
        "col2": ["x", "y", "y", "z", "z"],
        "color": [
            "lightgrey",
            "lightblue",
            "lightblue",
            "papayawhip",
            "papayawhip",
        ],
    }
    df = pl.DataFrame(data)
    return (df,)


@app.cell
def _(GT, df, html, mo):
    style_widget = mo.ui.slider(1, 6, label="Style Number")

    _colors = ["blue", "cyan", "pink", "green", "red", "gray"]
    color_widget = mo.ui.radio(
        options=_colors, value=_colors[0], label="Style Color", inline=True
    )

    row_striping_widget = mo.ui.switch(value=True, label="Row Striping?")

    gt = (
        GT(df)
        .tab_header(html(style_widget), html(color_widget))
        .tab_source_note(html(row_striping_widget))
        .opt_align_table_header("left")
    )
    return color_widget, gt, row_striping_widget, style_widget


@app.cell
def _(color_widget, gt, row_striping_widget, style_widget):
    gt.opt_stylize(
        style=style_widget.value,
        color=color_widget.value,
        add_row_striping=row_striping_widget.value,
    )
    return


if __name__ == "__main__":
    app.run()

marimo

Check out the full marimo code below or view it on molab.

import marimo

__generated_with = "0.13.15"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    import polars as pl
    from great_tables import GT, loc, style

    return GT, mo, pl


@app.cell
def _(pl):
    data = {
        "col1": [2, 5, 7, 10, 15],
        "col2": ["x", "y", "y", "z", "z"],
        "color": ["lightgrey", "lightblue", "lightblue", "papayawhip", "papayawhip"],
    }
    df = pl.DataFrame(data)
    return (df,)


@app.cell
def _(mo):
    style_widget = mo.ui.slider(1, 6, label="Select Style Number")
    mo.output.append(style_widget)

    _colors = ["blue", "cyan", "pink", "green", "red", "gray"]
    color_widget = mo.ui.radio(
        options=_colors, value=_colors[0], label="Select Style Color"
    )
    mo.output.append(color_widget)

    row_striping_widget = mo.ui.switch(value=True, label="Add Row Striping?")
    mo.output.append(row_striping_widget)
    return color_widget, row_striping_widget, style_widget


@app.cell
def _(GT, color_widget, df, row_striping_widget, style_widget):
    GT(df).opt_stylize(
        style=style_widget.value,
        color=color_widget.value,
        add_row_striping=row_striping_widget.value,
    )
    return


if __name__ == "__main__":
    app.run()

case_when()

The case_when() function accepts two arguments:

• caselist: A list of two-element tuples, where the first item is the condition (used in pl.when()), and the second is the corresponding result expression (used in .then()).
• otherwise: A fallback expression used in .otherwise() if no conditions match.

The given example demonstrates how case_when() can simplify conditional logic compared to the more verbose pl.when().then().otherwise() chain.

from functools import cache
from typing import Any

import polars as pl


def case_when(
    caselist: list[tuple[pl.Expr, pl.Expr]], otherwise: pl.Expr | None = None
) -> pl.Expr:
    """
    Simplifies conditional logic in Polars by chaining multiple `when-then` expressions.

    Parameters
    ----------
    caselist
        A list of (condition, value) pairs. Each condition is evaluated in order,
        and the corresponding value is returned when a condition is met.
    otherwise
        The fallback value to use if none of the conditions match.

    Returns
    -------
    pl.Expr

    Examples:
    -------
    ```python
    import polars as pl

    df = pl.DataFrame({"x": [1, 2, 3, 4]})

    expr = case_when(
        caselist=[
            (pl.col("x") < 2, pl.lit("small")),
            (pl.col("x") < 4, pl.lit("medium"))
        ],
        otherwise=pl.lit("large"),
    ).alias("size")

    # This is equivalent to writing:
    # expr = (
    #     pl.when(pl.col("x") < 2)
    #       .then(pl.lit("small"))
    #       .when(pl.col("x") < 4)
    #       .then(pl.lit("medium"))
    #       .otherwise(pl.lit("large"))
    #       .alias("size")
    # )

    df.with_columns(expr)
    ```
    shape: (4, 2)
    ┌─────┬────────┐
    │ x   ┆ size   │
    │ --- ┆ ---    │
    │ i64 ┆ str    │
    ├─────┼────────┤
    │ 1   ┆ small  │
    │ 2   ┆ medium │
    │ 3   ┆ medium │
    │ 4   ┆ large  │
    └─────┴────────┘
    """
    (first_when, first_then), *cases = caselist

    # first
    expr = pl.when(first_when).then(first_then)

    # middles
    for when, then in cases:
        expr = expr.when(when).then(then)

    # last
    expr = expr.otherwise(otherwise)

    return expr

Custom Expression Namespace

With case_when() in place, we can refactor the DiscreteSplitter expression namespace like this:

@cache
def _mod_expr(n: int) -> pl.Expr:
    return pl.int_range(pl.len(), dtype=pl.UInt32).mod(n)


def _litify(lits: list[Any]) -> list[pl.lit]:
    return [pl.lit(lit) for lit in lits]


@pl.api.register_expr_namespace("spt")
class DiscreteSplitter:
    def __init__(self, expr: pl.Expr) -> None:
        self._expr = expr

    def _get_expr(self, lits: list[Any], name: str):
        n = len(lits)
        mod_expr = _mod_expr(n)
        *litified, litified_otherwise = _litify(lits)
        caselist = [(mod_expr.eq(i), lit) for i, lit in enumerate(litified)]
        return case_when(caselist, litified_otherwise).alias(name)

    def binarize(self, lit1: Any, lit2: Any, name: str = "binarized") -> pl.Expr:
        return self.bucketize([lit1, lit2], name)

    def trinarize(
        self, lit1: Any, lit2: Any, lit3: Any, name: str = "trinarized"
    ) -> pl.Expr:
        return self.bucketize([lit1, lit2, lit3], name)

    def bucketize(self, lits: list[Any], name: str = "bucketized") -> pl.Expr:
        return self._get_expr(lits, name)

Now, bucketize() is the primary method that encapsulates the core logic for categorical mapping. binarize() and trinarize() are just convenient wrappers for common cases.

Here’s a simple example of using the custom expression namespace:

df = (
    pl.DataFrame({"n": [100, 50, 72, 83, 97, 42, 20, 51, 77]})
    .with_row_index(offset=1)
    .with_columns(
        pl.col("").spt.binarize("lightblue", "papayawhip"),
        pl.col("").spt.trinarize("one", "two", "three"),
        pl.col("").spt.bucketize([1, 2, 3, 4]),
    )
)

shape: (9, 5)
┌───────┬─────┬────────────┬────────────┬────────────┐
│ index ┆ n   ┆ binarized  ┆ trinarized ┆ bucketized │
│ ---   ┆ --- ┆ ---        ┆ ---        ┆ ---        │
│ u32   ┆ i64 ┆ str        ┆ str        ┆ i32        │
╞═══════╪═════╪════════════╪════════════╪════════════╡
│ 1     ┆ 100 ┆ lightblue  ┆ one        ┆ 1          │
│ 2     ┆ 50  ┆ papayawhip ┆ two        ┆ 2          │
│ 3     ┆ 72  ┆ lightblue  ┆ three      ┆ 3          │
│ 4     ┆ 83  ┆ papayawhip ┆ one        ┆ 4          │
│ 5     ┆ 97  ┆ lightblue  ┆ two        ┆ 1          │
│ 6     ┆ 42  ┆ papayawhip ┆ three      ┆ 2          │
│ 7     ┆ 20  ┆ lightblue  ┆ one        ┆ 3          │
│ 8     ┆ 51  ┆ papayawhip ┆ two        ┆ 4          │
│ 9     ┆ 77  ┆ lightblue  ┆ three      ┆ 1          │
└───────┴─────┴────────────┴────────────┴────────────┘

Custom DataFrame Namespace

Instead of relying on pl.DataFrame.with_row_index(), we can also use _mod_expr() directly to enable similar categorization.

Here’s how the DiscreteSplitter can be implemented as a custom DataFrame namespace:

@pl.api.register_dataframe_namespace("spt")
class DiscreteSplitter:
    def __init__(self, df: pl.DataFrame) -> None:
        self._df = df

    def _get_expr(self, lits: list[Any], name: str):
        n = len(lits)
        mod_expr = _mod_expr(n)
        *litified, litified_otherwise = _litify(lits)
        caselist = [(mod_expr.eq(i), lit) for i, lit in enumerate(litified)]
        return case_when(caselist, litified_otherwise).alias(name)

    def _get_final_df(self, lits: list[Any], name: str) -> pl.DataFrame:
        cls = type(self)
        expr = self._get_expr(lits, name)
        new_spt = cls(self._df.with_columns(expr))
        return new_spt._df

    def binarize(self, lit1: Any, lit2: Any, name: str = "binarized") -> pl.DataFrame:
        return self.bucketize([lit1, lit2], name=name)

    def trinarize(
        self, lit1: Any, lit2: Any, lit3: Any, name: str = "trinarized"
    ) -> pl.DataFrame:
        return self.bucketize([lit1, lit2, lit3], name=name)

    def bucketize(self, lits: list[Any], name: str = "bucketized") -> pl.DataFrame:
        return self._get_final_df(lits, name)

Example usage:

df = (
    pl.DataFrame({"n": [100, 50, 72, 83, 97, 42, 20, 51, 77]})
    .spt.binarize("lightblue", "papayawhip")
    .spt.trinarize("one", "two", "three")
    .spt.bucketize([1, 2, 3, 4])
    .with_row_index(offset=1)
)

shape: (9, 5)
┌───────┬─────┬────────────┬────────────┬────────────┐
│ index ┆ n   ┆ binarized  ┆ trinarized ┆ bucketized │
│ ---   ┆ --- ┆ ---        ┆ ---        ┆ ---        │
│ u32   ┆ i64 ┆ str        ┆ str        ┆ i32        │
╞═══════╪═════╪════════════╪════════════╪════════════╡
│ 1     ┆ 100 ┆ lightblue  ┆ one        ┆ 1          │
│ 2     ┆ 50  ┆ papayawhip ┆ two        ┆ 2          │
│ 3     ┆ 72  ┆ lightblue  ┆ three      ┆ 3          │
│ 4     ┆ 83  ┆ papayawhip ┆ one        ┆ 4          │
│ 5     ┆ 97  ┆ lightblue  ┆ two        ┆ 1          │
│ 6     ┆ 42  ┆ papayawhip ┆ three      ┆ 2          │
│ 7     ┆ 20  ┆ lightblue  ┆ one        ┆ 3          │
│ 8     ┆ 51  ┆ papayawhip ┆ two        ┆ 4          │
│ 9     ┆ 77  ┆ lightblue  ┆ three      ┆ 1          │
└───────┴─────┴────────────┴────────────┴────────────┘

Conclusion

Extracting the conditional logic into a standalone case_when() function turned out to be both a practical and satisfying exercise—perfect for a rainy afternoon of coding. It not only improves readability but also makes the branching logic easier to reuse and reason about.