Skip to content

Make interpreters compatible with Python type-checking #355

@eb8680

Description

@eb8680

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions