"Pure" Type-Constrained Package Sorter

The Problem

Objective

Imagine you work in a robotic automation factory, and your objective is to write a function for one of its robotic arms that will dispatch the packages to the correct stack according to their volume and mass.

Rules

Sort the packages using the following criteria:

• A package is bulky if its volume (Width x Height x Length) is greater than or equal to 1,000,000 cm³ or when one of its dimensions is greater or equal to 150 cm.
• A package is heavy when its mass is greater or equal to 20 kg.

You must dispatch the packages in the following stacks:

• STANDARD: standard packages (those that are not bulky or heavy) can be handled normally.
• SPECIAL: packages that are either heavy or bulky can't be handled automatically.
• REJECTED: packages that are both heavy and bulky are rejected.

Goal:

Implement the function sort(width, height, length, mass) (units are centimeters for the dimensions and kilogram for the mass). This function must return a string: the name of the stack where the package should go.

Thoughts

Instead of using a typical imperative approach, I opted for a more declarative, type-driven design using pydantic. The core idea was to encode the sorting logic directly into the type system by defining models that represent each package category (STANDARD, SPECIAL, and REJECTED) with appropriate constraints. By casting the input data into a PackageConsumer model, pydantic handles the validation and selects the correct subtype (PackageT) based on the input's properties. This allows the classification logic to be enforced and executed through type validation alone. As a result, determining the appropriate stack is simply a matter of reading the label attribute of the returned model, cleanly separating the logic from control flow and making the code both expressive and easy to extend.

Installation/Requirements

I'm running Python v3.13.3, but this should work with versions 3.8+.

Dependencies

Only pydantic is required, numpy & matplotlib are both mainly used for testing/visualization purposes.

Setup Virtual Environment

I generally use pipenv or uv, but the setup process will remain the same.

Using pipenv:

pipenv install
pipenv shell

Using uv:

uv venv
uv pip install -r requirements.txt
source .venv/bin/activate  # or use `uv venv --python` to specify a Python version

Running Locally

In a shell or IDE, the base solution.py script provides testing and visualization functions.

Plotting the Problem Space

The command below renders a 2D visualization representing a slice of a 4D result space (varying width, height, length, and mass). Each grid cell represents a 2D slice (Width x Height) of the original 4D array for a specific combination of Length (L) and Mass (M). This "stitched" layout enables you to view all (L, M) combinations simultaneously in one figure, rather than inspecting each subplot separately.

python3 -c "import solution; solution.viz()"

The result should look like this:

Tests & Exploration

These aren't formal tests (I've omitted pytest to keep dependencies minimal), but the command below should run without errors.

python3 -c "import solution; solution.run_all_tests()"

To do more interesting testing/inspecting of the data follow these steps:

Load up a REPL:

python3 -i solution.py # or ipython3

Run a test:

>>> results = simple_test()
>>> results, dims, masses = test_sorting_on_meshgrid()
>>> failing_test1()

Visualize the problem space:

>>> viz()