pythondesign-patternsenumsdesign-principles

Best approach to map enums to functions


I have 4 distinct file types. Each of these file types maps to a different file name pattern given an index. What is the best way to do this? Here is one approach:

from pathlib import Path
from enum import Enum

base_path = "/a/b/c/d"

class FileType(Enum):
    IMG = lambda index: f"img_{index}.nii.gz"
    IMG_SEGMENTATION = lambda index: f"img_{index}_segmentation.nii.gz"
    SUB_IMG = lambda index: f"img_sub_{index}.nii.gz"
    SUB_IMG_SEGMENTATION = lambda index: f"img_sub_{index}_segmentation.nii.gz"

def get_path(file_type: FileType, index):
    path = Path(base_path) / file_type.value(index)
    if not path.exists():
        raise FileNotFoundError
    return path

The problem that I don't like with the above approach is that it uses an enum to store a function, not an int, which I think goes against the spirit of what an enum is for.

An alternative is the following, adding an intermediate function to determine the associated function. I don't like the following either because if a file type is added to enum, it would also need to be added to file_type_to_name.

from pathlib import Path
from enum import Enum, auto

base_path = "/a/b/c/d"

class FileType(Enum):
    IMG = auto()
    IMG_SEGMENTATION = auto()
    SUB_IMG = auto()
    SUB_IMG_SEGMENTATION = auto()

def file_type_to_name(file_type: FileType, index: int) -> str:
    if file_type == FileType.IMG:
        return f"img_{index}.nii.gz"
    if file_type == FileType.IMG_SEGMENTATION:
        return f"img_{index}_segmentation.nii.gz"
    if file_type == FileType.SUB_IMG:
        return f"img_sub_{index}.nii.gz"
    if file_type == FileType.SUB_IMG_SEGMENTATION:
        return f"img_sub_{index}_segmentation.nii.gz"

def get_path(file_type: FileType, index):
    path = Path(base_path) / file_type_to_name(file_type, index)
    if not path.exists():
        raise FileNotFoundError
    return path

Solution

  • How about using a dictionary to map each FileType to a function? This is cleaner than the if/else solution and also allows handling various error more elegantly. Another benefit is that it decouples the path-finding logic from the enum itself, which could be beneficial in some cases. For example:

    class FileType(Enum):
        IMG = auto()
        IMG_SEGMENTATION = auto()
        SUB_IMG = auto()
        SUB_IMG_SEGMENTATION = auto()
    
    FILE_TYPE_TO_PATH: dict[FileType, Callable[[int], str]] = {
        FileType.IMG: lambda index: f"img_{index}.nii.gz",
        FileType.IMG_SEGMENTATION: lambda index: f"img_{index}_segmentation.nii.gz",
        FileType.SUB_IMG: lambda index: f"img_sub_{index}.nii.gz",
        FileType.SUB_IMG_SEGMENTATION: lambda index: f"img_sub_{index}_segmentation.nii.gz",
    }
    
    def file_type_to_name(file_type: FileType, index: int) -> str:
        try:
            return FILE_TYPE_TO_PATH[file_type](index)
        except KeyError:
            ## handle the fact that a FileType does not have a matching function
    
    def file_type_to_name(file_type: FileType, index: int) -> str:
        path_function = FILE_TYPE_TO_PATH.get(file_type, some_default_function)
        return path_function(index)
    
    def file_type_to_name(file_type: FileType, index: int) -> str:
        if (path_function := FILE_TYPE_TO_PATH.get(file_type)):
            return path_function(index)
        raise SomeHorribleError("Oh no")