solution.py

from pydantic import BaseModel, Field, computed_field
from enum import StrEnum, IntEnum
from typing import Literal

"""
Assumptions:
------------

All values in this script assume dimension units are:
    - w/h/l : cm
    - mass  : kg

Any values that are negative or zero will be automatically
rejected and will result in a pydantic.ValidationError
"""

class Constraints:
    """
    Defines hard-coded constraints for the given problem:
        - A package is bulky if volume >= 1,000,000 cm^3
            * A package can also be bulky if a dimension >= 150 cm
        - A package is heavy if mass >= 20kg
    """
    MAX_STANDARD_VOLUME           : float = 1e6   # cm^3
    MAX_STANDARD_DIMENSION_LENGTH : float = 150.0 # cm
    MAX_STANDARD_MASS             : float = 20.0  # kg

class VolumeLabel(StrEnum):
    STANDARD = 'STANDARD'
    BULKY    = 'BULKY'

class MassLabel(StrEnum):
    STANDARD = 'STANDARD'
    HEAVY    = 'HEAVY'

class PackageLabel(IntEnum):
    STANDARD = 1
    SPECIAL  = 2
    REJECTED = 3

class Volume(BaseModel):
    value : float = Field(..., gt=0.)
    label : VolumeLabel

class StandardVolume(Volume):
    value : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_VOLUME
    )
    label : VolumeLabel = VolumeLabel.STANDARD

class BulkyVolume(Volume):
    label : VolumeLabel = VolumeLabel.BULKY

VolumeT = (
    StandardVolume | BulkyVolume
)

class VolumeConsumer(BaseModel):
    volume : VolumeT

class Dimension(BaseModel):
    width  : float = Field(..., gt=0.)
    height : float = Field(..., gt=0.)
    length : float = Field(..., gt=0.)

    @computed_field
    @property
    def volume(self) -> VolumeT:
        return VolumeConsumer(
            volume={
                'value': (
                    self.width
                    * self.height
                    * self.length
                )
            }
        ).volume

class StandardDimension(Dimension):
    width : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_DIMENSION_LENGTH
    )
    height : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_DIMENSION_LENGTH
    )
    length : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_DIMENSION_LENGTH
    )

class BulkyDimension(Dimension):
    ...

DimensionT = (
    StandardDimension | BulkyDimension
)

class DimensionConsumer(BaseModel):
    dimension : DimensionT

    @classmethod
    def from_args(
        cls, *
        , width  : int | float
        , height : int | float
        , length : int | float
    ) -> DimensionT:
        return cls(
            dimension={
                'width': width
                , 'height': height
                , 'length': length
            }
        ).dimension

class Package(BaseModel):
    dimension : DimensionT
    volume    : VolumeT
    mass      : float
    label     : PackageLabel

class StandardPackage(Package):
    dimension : StandardDimension
    volume    : StandardVolume
    mass      : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_MASS
    )
    label     : PackageLabel = PackageLabel.STANDARD

class BulkyPackage(Package):
    dimension : DimensionT
    volume    : VolumeT
    mass      : float = Field(
        ...
        , gt=0.
        , lt=Constraints.MAX_STANDARD_MASS
    )
    label     : PackageLabel = PackageLabel.SPECIAL

class HeavyPackage(StandardPackage):
    """
    Note: we're using StandardPackage to avoid
    re-writing the same dimension/volume types again,
    since a heavy package is only heavy when having
    standard dimensions/volume (otherwise it's considered rejected)
    """
    mass  : float = Field(..., ge=Constraints.MAX_STANDARD_MASS)
    label : PackageLabel = PackageLabel.SPECIAL

class RejectedPackage(Package):
    dimension : DimensionT
    volume    : VolumeT
    mass      : float = Field(..., ge=Constraints.MAX_STANDARD_MASS)
    label     : PackageLabel = PackageLabel.REJECTED

# Define the generic Package Type, in this problem domain
# a package can be one of the following:
#   - Standard
#   - Heavy -> Special
#   - Bulky -> Special
#   - Rejected
PackageT = (
    StandardPackage
    | HeavyPackage
    | BulkyPackage
    | RejectedPackage
)

class PackageConsumer(BaseModel):
    package : PackageT

    @classmethod
    def from_args(
        cls, *
        , mass : int | float
        , **dimensions
    ) -> PackageT:
        dim = DimensionConsumer.from_args(
            **dimensions
        )
        return cls(package={
            'dimension': dim
            , 'volume': dim.volume
            , 'mass': mass
        }).package

"""
Implementation of sort just involves casting the arguments into a PackageT
the optional ret_type parameter allows caller control of how
the result is returned, with options to return:
  - label name (the enum name)
  - label value (the value of the enum, int for plotting purposes)
  - label enum
  - model (PackageT)

The default selection is name to adhere to the takehome requirements
"""
def sort(
    width    : int | float
    , height : int | float
    , length : int | float
    , mass   : int | float
    , ret_type : str = 'name'
) -> PackageT | PackageLabel | str | int:
    pkg = PackageConsumer.from_args(
        width=width
        , height=height
        , length=length
        , mass=mass
    )

    match ret_type:
        case 'label':
            return pkg.label # PackageLabel
        case 'pkg' | 'package':
            return pkg # PackageT
        case 'label value' | 'value' | 'lblvalue':
            return pkg.label.value # int
        case 'label name' | 'name' | 'lblname':
            return pkg.label.name # str
        case _:
            raise ValueError(f"{re_type} is not a valid return type!")

# --------------- TESTING ---------------- #

from itertools import product
import matplotlib.pyplot as plt
import numpy as np

def simple_test():
    """
    Creates the full combination space (or product) of
    dims x dims x dims x masses and runs the sort operation

    """
    dims   = [10, 50, 100, 150]
    masses = [5, 10, 20, 25]

    results = {}

    for w, h, l, m in product(dims, dims, dims, masses):
        results[(w, h, l, m)] = sort(w, h, l, m, 'package')

    return results

def failing_test1():
    """
    Raises ValidationError due to zero width
    """
    sort(0, 10, 10, 1)

def failing_test2():
    """
    Raises ValidationError due to zero width and negative height
    """
    sort(0, -1, 10, 1)

def sort_for_meshgrid(
    w, h, l, m
):
    """
    Wrapper function of sort, required for meshgrid testing
    sets the `ret_type` to value to produce an int result
    """
    return sort(w, h, l, m, 'value')

def test_sorting_on_meshgrid():
    """
    Creates a 4D meshgrid of shape (dims, dims, dims, masses)
    and runs a "vectorized" version of sort to produce a
    4D result array. Also returns the dimensions and masses
    to be used for plotting purposes

    The dimensions width, length, height (w/l/h) are chosen
    arbitrarily since we are re-using `dims` across 3D.
    """
    dims   = np.array([1, 5, 50, 100, 150, 200, 1000])
    masses = np.array([1, 5, 10, 15, 20, 25, 50])

    W, L, H, M = np.meshgrid(dims, dims, dims, masses)
    vectorized = np.vectorize(sort_for_meshgrid)

    return vectorized(W, L, H, M), dims, masses

def viz():
    """
    Stitches together multiple 4D -> 2D sub-arrays
    into one giant 2D array to produce a giant
    grid plot of results. Each grid item in the
    rendered plot is a "patch" of the WxH result sub-array
    indexed at a given length and mass. This lets us view
    multiple slices of the 4D array in one figure, without
    having to inspect multiple subplots.

    Hard-coded to currently vary the length & mass
    over height & width. Since l/w/h are all the same,
    this rendered slice of the full 4D result array
    is enough to showcase the problem domain for a wide
    variety of inputs.
    """
    result, dims, masses = test_sorting_on_meshgrid()

    n_w, n_l, n_h, n_m = result.shape

    # Flatten the data by stitching slices over W × H
    # for each (L, M) combination
    # Resulting shape will be (n_w * n_l, n_h * n_m)
    stitched = np.zeros((n_w * n_l, n_h * n_m))

    for i_l in range(n_l):
        for i_m in range(n_m):
            slice_wh = result[:, i_l, :, i_m]

            row_start = i_l * n_w
            row_end   = (i_l + 1) * n_w

            col_start = i_m * n_h
            col_end   = (i_m + 1) * n_h

            stitched[row_start:row_end, col_start:col_end] = slice_wh

    fig, ax = plt.subplots(figsize=(10, 8))
    im = ax.imshow(stitched, cmap='viridis', origin='lower', aspect='auto')

    for i in range(1, n_l):
        ax.axhline(i * n_w - 0.5, color='white', linewidth=0.5)
    for j in range(1, n_m):
        ax.axvline(j * n_h - 0.5, color='white', linewidth=0.5)

    ax.set_xticks([i * n_h + n_h / 2 for i in range(n_m)])
    ax.set_xticklabels([f"M={m}" for m in masses], rotation=45)

    ax.set_yticks([i * n_w + n_w / 2 for i in range(n_l)])
    ax.set_yticklabels([f"L={l}" for l in dims])

    ax.set_title("Package Sorting Results: WxH slices arranged by LxM")
    fig.colorbar(im, ax=ax, label='Result Value')
    plt.tight_layout()
    plt.show()

def run_all_tests():
    def run_failed_test(func):
        try:
            func()
        except Exception as e:
            return e

    results = [
        ('simple_test', simple_test())
        , ('sort_for_meshgrid', test_sorting_on_meshgrid())
    ]

    failures = [
        ('failing_test1', run_failed_test(failing_test1))
        , ('failing_test2', run_failed_test(failing_test2))
    ]

    return results, failures