SQLAlchemy has a cool feature where I can write an expression, and depending on the input type, the same expression can build a SQL condition, or evaluate to a value. I’m trying to do the same in my code, using functions instead of expressions, but I can’t figure out how to annotate my functions for pypy to correctly check the types.

The below code works, both assertions pass. However, mypy complains:

example.py:13: error: Argument 1 has incompatible type "SQLColumnExpression[Any]"; expected "str | int | bool | date | datetime"  [arg-type]
example.py:14: error: Argument 1 has incompatible type "str | int | bool | date | datetime"; expected "SQLColumnExpression[Any]"  [arg-type]

As I understand it, the problem is that the var argument can be either of the callables in the Union, but the code expects it to be both. How would I declare such a type, or better, can I find something somewhere in SQLAlchemy to import ?

from sqlalchemy import SQLColumnExpression, Column, Integer, SQLColumnExpression
from sqlalchemy.orm import declarative_base
from datetime import date, datetime
from typing import Union, Callable, Any, overload

Value = Union[str, int, bool, date, datetime]
Operator = Union[
    Callable[[SQLColumnExpression[Any], Value], SQLColumnExpression[Any]],
    Callable[[Value, Value], bool],
]

def some_function(var: Operator, column: SQLColumnExpression[Any], value: Value) -> None:
    assert isinstance(var(column, value), SQLColumnExpression)
    assert isinstance(var(value, value), bool)

@overload
def op(column: SQLColumnExpression[Any], value: Value) -> SQLColumnExpression[Any]: ...

@overload
def op(column: Value, value: Value) -> bool: ...

def op(column, value):  # type: ignore
    print("OP", type(column), type(value))
    return column == 0

Base = declarative_base()
class Model(Base):  # type: ignore
    __tablename__ = "example"
    id = Column(Integer, primary_key=True)

some_function(op, Model.id, 0)