How to type dynamically created classes so mypy can lint them properly

329 Views Asked by At

I'm looking to refactor the function task decorator of a dataflow engine I contribute to called Pydra so that the argument types can be linted with mypy. The code captures the arguments to be passed to the function at workflow construction time and stores them in a dynamically created Inputs class (designed using python-attrs), in order to delay the execution of the function until runtime of the workflow.

The following code works fine, but mypy doesn't know the types of the attributes in the dynamically generated classes. Is there a way to specify them dynamically too

import typing as ty
import attrs
import inspect
from functools import wraps


@attrs.define(kw_only=True, slots=False)
class FunctionTask:

    name: str
    inputs: ty.Type[ty.Any]  # Use typing.Type hint for inputs attribute

    def __init__(self, name: str, **kwargs):
        self.name = name
        self.inputs = type(self).Inputs(**kwargs)


def task(function: ty.Callable) -> ty.Type[ty.Any]:

    sig = inspect.signature(function)

    inputs_dct = {
        p.name: attrs.field(default=p.default) for p in sig.parameters.values()
    }
    inputs_dct["__annotations__"] = {
        p.name: p.annotation for p in sig.parameters.values()
    }

    @wraps(function, updated=())
    @attrs.define(kw_only=True, slots=True, init=False)
    class Task(FunctionTask):
        func = staticmethod(function)

        Inputs = attrs.define(type("Inputs", (), inputs_dct))  # type: ty.Any

        inputs: Inputs = attrs.field()

        def __call__(self):
            return self.func(
                **{
                    f.name: getattr(self.inputs, f.name)
                    for f in attrs.fields(self.Inputs)
                }
            )

    return Task

which you can use with

@task
def myfunc(x: int, y: int) -> int:
    return x + y

# Would be great if `x` and `y` arguments could be type-checked as ints
mytask = myfunc(name="mytask", x=1, y=2)

# Would like mypy to know that mytask has an inputs attribute and that it has an int
# attribute 'x', so the linter picks up the incorrect assignment below
mytask.inputs.x = "bad-value"

mytask()

I would like mypy to know that mytask has an inputs attribute and that it has an int attribute 'x', so the linter picks up the incorrect assignment of "bad-value". If possible, it would be great if the keyword args to the myfunc.__init__ are also type-checked.

Is this possible? Any tips on things to try?

EDIT: To try to make what I'm trying to do a bit clearer, here is an example of what one of the dynamically generated classes would look like if it was written statically

@attrs.define
class StaticTask:

    @attrs.define
    class Inputs:
        x: int
        y: int

    name: str
    inputs: Inputs = attrs.field()

    @staticmethod
    def func(x: int, y: int) -> int:
        return x + y

    def __init__(self, name: str, x: int, y: int):
        self.name = name
        self.inputs.x = x
        self.inputs.y = y

    def __call__(self):
        return self.func(x=self.inputs.x, y=self.inputs.y)

In this case

mytask2 = StaticTask(name="mytask", x=1, y=2)

mytask2.inputs.x = "bad-value"

the final line gets flagged by mypy as setting a string to an int field. This is what I would like my dynamically created classes to replicate.

2

There are 2 best solutions below

1
chepner On

I may be misunderstanding what you need, but if all you are trying to do is add an additional typed parameter to a function, you can do that using typing.ParamSpec.

P = ty.ParamSpec('P')
RV = ty.TypeVar('RV')


def task(f: ty.Callable[P, RV]) -> ty.Callable[ty.Concatenate[str, P], RV]:
    def _(name: str, **kwargs: P.kwargs) -> RV:
        # do something with name?
        return f(**kwargs)
    return _
0
Tom Close On

Reading through the mypy docs more closely, https://mypy.readthedocs.io/en/stable/dynamic_typing.html, it looks like what I wanted to be able to do isn't possible and the best we can do is

def task(function: ty.Callable) -> ty.Type[FunctionTask]:

@chepner's answer could get you most of the way there if you didn't need to access the inputs/outputs directly and just needed to change the signature to include the name arg