Ability to use Callable type alias when annotating functions

[OlegAlexander]

Hello. I would like to revive the discussion about Callable type aliases. The original discussion is here, and the related discussion about @declared_type is here and here.

The original thread contains a great example of the problem and a proposed solution:

from typing import Callable
import math

ActivationFunction = Callable[[float], float]

sigmoid: ActivationFunction
def sigmoid(x):
     return 1 / (1 + math.exp(-x))

relu: ActivationFunction
def relu(x):
     return max(0, x)

Another prosed solution is to use the @declared_type decorator:

@declared_type(ActivationFunction)
def sigmoid(x):
     return 1 / (1 + math.exp(-x))

@declared_type(ActivationFunction)
def relu(x):
     return max(0, x)

Currently, this can be done with Protocols, but it's too verbose:

from typing import Protocol
import math

class ActivationFunction(Protocol):
    def __call__(self, x: float) -> float:
        ...

class Sigmoid(ActivationFunction):
    def __call__(self, x: float) -> float:
        return 1 / (1 + math.exp(-x))

class Relu(ActivationFunction):
    def __call__(self, x: float) -> float:
        return max(0, x)

sigmoid = Sigmoid()
relu = Relu()

print(sigmoid(0.5))
print(relu(0.5))

If the function is a single expression, it can be done very nicely with lambdas:

from typing import Callable
import math

ActivationFunction = Callable[[float], float]

sigmoid: ActivationFunction = lambda x: 1 / (1 + math.exp(-x))
relu: ActivationFunction = lambda x: max(0, x)

print(sigmoid(0.5))
print(relu(0.5))

Here's the equivalent FSharp code, with the only difference being that FSharp supports multiline lambdas:

type ActivationFunction = float -> float

let sigmoid: ActivationFunction = fun x -> 1.0 / (1.0 + exp(-x))
let relu: ActivationFunction = fun x -> max 0.0 x

printfn "%f" (sigmoid 0.5)
printfn "%f" (relu 0.5)

I'm asking about this feature because I'm reading a book called Domain Modeling Made Functional by Scott Wlaschin. In this book, the author recommends modeling the domain with types, including function type aliases, and focusing on the implementation details later. Here's an example from Chapter 9:

I'd like to try a similar approach in Python. Has there been any further discussion about this topic, or is there a workaround I'm unaware of? Thank you!

[erictraut]

I'm not clear on the problem you're trying to solve. Are your needs not addressed by standard function type annotations?

def sigmoid(x: float) -> float:
    return 1 / (1 + math.exp(-x))

def relu(x: float) -> float:
    return max(0, x)

print(sigmoid(0.5))
print(relu(0.5))

sigmoid("bad arg") # Type checker error - in correct argument type

A protocol is useful if you want to create another interface that accepts any function that accepts a single float parameter named x and produces a float result.

class ActivationFunction(Protocol):
    def __call__(self, x: float) -> float:
        ...

def invoke_activation(val: float, activation: ActivationFunction) -> float:
    return activation(val)

invoke_activation(0.5, sigmoid)
invoke_activation(0.5, relu)
invoke_activation(0.5, lambda x: math.tanh(x))

If the parameter name isn't important (i.e. you plan to invoke it only by a positional argument), you can use a simple Callable[[float], float] annotation rather than a protocol.

ActivationFunction = Callable[[float], float]

[OlegAlexander]

The approach recommended in the book is to capture as much domain logic in the type system as possible. For example:

# domain.py

from typing import NewType, Callable

# Email validation workflow
UnvalidatedEmail = NewType('UnvalidatedEmail', str)
ValidEmail = NewType('ValidEmail', str)
EmailValidationError = NewType('EmailValidationError', str)
ValidateEmail = Callable[[UnvalidatedEmail], ValidEmail | EmailValidationError] 

The author argues that domain.py is so transparent that even a non-programmer client can understand it. This file also serves as "executable documentation."

Later, when we implement the validate_email function, how do we tell mypy it should have type ValidateEmail? I've found a satisfactory way, but please let me know if my solution can be improved upon:

# email_workflow.py

from domain import UnvalidatedEmail, ValidEmail, EmailValidationError, ValidateEmail

def _validate_email(email: UnvalidatedEmail) -> ValidEmail | EmailValidationError:
    if '@' not in email:
        return EmailValidationError(f'Invalid email: {email}')
    return ValidEmail(email)

validate_email: ValidateEmail = _validate_email

While I'm mostly satisfied with this solution, I still think the intent can be made much clearer with something like this:

# email_workflow.py

from domain import UnvalidatedEmail, ValidEmail, EmailValidationError, ValidateEmail

@declared_type(ValidateEmail)
def validate_email(email: UnvalidatedEmail) -> ValidEmail | EmailValidationError:
    if '@' not in email:
        return EmailValidationError(f'Invalid email: {email}')
    return ValidEmail(email)

Possible mypy bug

As an aside, while experimenting with my workaround, I found a possible bug in mypy.

Let's start with a working baseline:

from typing import Callable

ActivationFunction = Callable[[float], float]

def _relu(x: float) -> float:
    return max(0, x)

relu: ActivationFunction = _relu

print(relu(0.5))

Now let's rename _relu to relu and reassign it to relu:

from typing import Callable

ActivationFunction = Callable[[float], float]

def relu(x: float) -> float:
    return max(0, x)

relu: ActivationFunction = relu # I expected a no-redefinition error here

print(relu(0.5))

More disturbingly, let's change the function signature of relu:

from typing import Callable

ActivationFunction = Callable[[float], float]

def relu(x: float, y: float) -> float:
    return max(0, x)

relu: ActivationFunction = relu # I definitely expected a type error here

print(relu(0.5)) # But the type error is here

On the other hand, if we rename relu to add, the type error occurs where expected:

from typing import Callable

ActivationFunction = Callable[[float], float]

def add(x: float, y: float) -> float:
    return x + y

relu: ActivationFunction = add # Type error is here as expected

print(relu(0.5)) 

As you can see, it has something to do with the function name being the same during definition and reassignment.

Is this a bug or intended behavior? If it's a bug, I can report it on the mypy github.

[OlegAlexander]

@ilevkivskyi, I'm so sorry to tag you personally on this, but I think it's justified because you were the original champion of @declared_type here: python/mypy#1641 (comment) Has there been any progress on @declared_type since then? Do you still think it's worth adding? Thank you.