Search code examples
pythonpython-3.xtype-hintingmypy

What's a good pattern for typehinting with `Literal` and then validating at runtime?


Let's say I have a class:

from typing import Literal


class Computer:
    def __init__(self, operation: Literal['floor', 'ceil', 'square']):
        if operation not in ('floor', 'ceil', 'square'):
            raise ValueError(f'Invalid operation "{operation}"')
        self._operation = operation

    # ... lots of methods that use `self._operation`, e.g.
    def is_much_larger(self, value: float) -> bool:
        if self._operation == 'square':
            if value > 1:
                return True
            else:
                return False
        else:
            return False

I would like to maintain the list of viable values, i.e. 'floor', 'ceil', 'square', in only one place. What's a good recipe for this?

The MyPy docs show some ideas with this suggestion of assert_never but that's just for MyPy to give an error before runtime, not as a way to maintain a single list of values.


Solution

  • You can introspect those arguments using the typing module.

    >>> import typing
    >>> Operation = typing.Literal['floor', 'ceil', 'square']
    >>> typing.get_args(Operation)
    ('floor', 'ceil', 'square')
    

    You could use a type alias then get the values you want using typing.get_args, so something like:

    import typing
    
    # in Python >= 3.10 , use explicit type alias:
    # Operation: typing.TypeAlias = typing.Literal['floor', 'ceil', 'square'
    Operation = typing.Literal['floor', 'ceil', 'square']
    
    
    class Computer:
        def __init__(self, operation: Operation:
            if operation not in typing.get_args(Operation):
                raise ValueError(f'Invalid operation "{operation}"')
            self._operation = operation
    
        # ... lots of methods that use `self._operation`, e.g.
        def is_much_larger(self, value: float) -> bool:
            if self._operation == 'square':
                if value > 1:
                    return True
                else:
                    return False
            else:
                return False