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.
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.