Currently, while Funsor makes heavy use of multipledispatch on parametric types, interpretations and Funsor terms are not compatible with static type-checkers like MyPy or even runtime type-checkers like pytypes. It would improve the development experience, increase performance and beneficially constrain Funsor's design to make Funsor interpreters and lazy expressions more compatible with Python's standard typing module and even with MyPy.
At a high level, this seems straightforward: just make all subclasses of funsor.Funsor inherit from typing.Generic and move type annotations from register calls to function signatures, e.g.:
Op = TypeVar("Op", bound=ops.AssociativeOp)
Arg = TypeVar("Arg", bound=Funsor)
ReducedVars = TypeVar("ReducedVars", bound=FrozenSet[str])
class Reduce(Funsor, Generic[Op, Arg, ReducedVars]):
def __init__(self, op: Op, arg: Arg, reduced_vars: ReducedVars):
... # note we should be able to remove all hand-written type assertions here
Lhs = TypeVar("Lhs", bound=Funsor)
Rhs = TypeVar("Rhs", bound=Funsor)
class Binary(Funsor, Generic[Op, Lhs, Rhs]):
def __init__(self, op, lhs, rhs):
...
# a random example pattern (a+b).sum() -> (a.sum() + b.sum())
@eager.register(Reduce)
def eager_reduce_example_pattern(
op: ops.AddOp,
arg: Binary[ops.AddOp, Funsor, Funsor],
reduced_vars: FrozenSet[str]
):
return Binary(arg.op, Reduce(op, arg.lhs, reduced_vars), Reduce(op, arg.rhs, reduced_vars))
We'd also have to modify KeyedRegistry along the lines of mrocklin/multipledispatch#69 using pytypes.deep_type for dispatch to continue to work correctly, and define FunsorMeta.__subclasscheck__ with pytypes.is_subtype rather than the custom logic it currently contains.
Unfortunately, there are lots of design details that may prevent fully statically checkable interpreters or at least make them undesirable for other reasons. For example, the obtrusive type variables in the snippet would seem to be mandatory - if their creation were automated and folded into FunsorMeta.__init__, the result would not be compatible with static type checkers like MyPy, which seems to be incompatible with any internal logic in metaclasses. It's also not clear that the way we construct finer types at runtime with reflect between intepreter calls would be compatible with MyPy.
These caveats may not apply to runtime type checkers like pytypes, which could still cut out lots of boilerplate and improve performance, so perhaps this is the best starting point. Note that there are Python version compatibility issues upstream in pytypes that might cause other problems (but seem to be mostly resolved in the master branch).
Note that this issue is orthogonal to #351, which is about the types of Funsors themselves rather than the type signatures of interpretations, despite the similar motivations.
Currently, while Funsor makes heavy use of
multipledispatchon parametric types, interpretations and Funsor terms are not compatible with static type-checkers like MyPy or even runtime type-checkers like pytypes. It would improve the development experience, increase performance and beneficially constrain Funsor's design to make Funsor interpreters and lazy expressions more compatible with Python's standardtypingmodule and even with MyPy.At a high level, this seems straightforward: just make all subclasses of
funsor.Funsorinherit fromtyping.Genericand move type annotations fromregistercalls to function signatures, e.g.:We'd also have to modify
KeyedRegistryalong the lines of mrocklin/multipledispatch#69 usingpytypes.deep_typefor dispatch to continue to work correctly, and defineFunsorMeta.__subclasscheck__withpytypes.is_subtyperather than the custom logic it currently contains.Unfortunately, there are lots of design details that may prevent fully statically checkable interpreters or at least make them undesirable for other reasons. For example, the obtrusive type variables in the snippet would seem to be mandatory - if their creation were automated and folded into
FunsorMeta.__init__, the result would not be compatible with static type checkers like MyPy, which seems to be incompatible with any internal logic in metaclasses. It's also not clear that the way we construct finer types at runtime withreflectbetween intepreter calls would be compatible with MyPy.These caveats may not apply to runtime type checkers like pytypes, which could still cut out lots of boilerplate and improve performance, so perhaps this is the best starting point. Note that there are Python version compatibility issues upstream in pytypes that might cause other problems (but seem to be mostly resolved in the master branch).
Note that this issue is orthogonal to #351, which is about the types of Funsors themselves rather than the type signatures of interpretations, despite the similar motivations.