Search code examples
pythontype-hinting

Giving names to returned values with python type hints


Let's say I have a function like

def basic_stats(data: List[float]):
    """This returns a tuple with the mean and median values of the data
    """
    return (np.mean(data), np.median(data))

and I want to use type hints to self-document the code instead of a comment. The standard way isn't super clear:

def basic_stats(data: List[float]) -> Tuple[float, float]:

because you can't tell by looking at the hint which output float is which.

Sometimes I've seen people do this:

def basic_stats(data: List[float]) -> Tuple["mean", "median"]:

which compiles just fine, and actually gives the humans what they need, but doesn't help the IDEs or mypys. (Although I'm yet to see those actually be useful in the real world yet, but that's maybe a different story.)

Is there any way to do both which is valid, pythonic, and useful? I'd like to be able to do something like this:

def basic_stats(data: List[float]) -> Tuple[mean:float, median:float]:

or should we just stick with something like

def basic_stats(data: List[float]) -> Tuple["mean:float", "median:float"]:

Or anything coming in a proposed PEP to solve this?


Solution

  • The way I normally do this—assigning new names to float—may not be to everyone's liking, but mypy accepts it and it makes the code completion so much more helpefull. It can also reduce errors, such as when computing a tuple (shape: float, scale: float) or similar where mixing the return value up would be really bad, but not a "code error" as such since they are both floats.

    Here's what I would have done with your code:

    from typing import List, Tuple
    
    
    mean = float
    median = float
    
    
    def basic_stats(data: List[float]) -> Tuple[mean, median]:
        """
        Compute the mean and median of a list of numbers
        """
        N = len(data)
        assert N > 0
        mean: float = sum(data) / N
        median: float = sorted(data)[N//2]
        return mean, median
        
        
    print(basic_stats([1.0, 1.0, 1.0]))
    print(basic_stats([3.0, 1.0, 2.0]))