CODEWITHSUNDEEP
pythonpython-typingmypyluigi
Achim
Achim

Reputation: 15692

Type annotations for parameters of Luigi tasks

I'm using Luigi in my Python project, so I have classes that look like this:

class MyTask(luigi.Task):
    my_attribute = luigi.IntParameter()

I would like to add a type annotation to my_attribute so that mypy will be aware that it is an integer. Or rather "will be an integer", because obviously it is not yet. It will become an integer due to "metaclass magic":

t = MyTask(my_attribute=5)
print(t.my_attribute) # <- t.my_attribute is an int, not an IntParameter

What's the proper way to annotate this attribute? Is it possible at all? I'm just a Luigi user and not a maintainer or contributor, so changing Luigi is not an option. At least not short term.

Upvotes: 0

Views: 112

Answers (1)

InSync
InSync

Reputation: 10437

This problem was raised in #2542, which was created in 2018 and automatically closed in 2019 as "stale". This perhaps means that the maintainers of luigi are not interested in adding proper type hints or creating a Mypy plugin for it.

A workaround would be to lie that IntParameter is a descriptor:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    class IntParameter:
        def __get__(self, instance: Any, owner: type[Any] | None, /) -> int: ...
else:
    from luigi import IntParameter

class MyTask(luigi.Task):
    my_attribute = IntParameter()

t = MyTask(my_attribute = 5)
reveal_type(t.my_attribute)  # int

If you are using other magic classes of the same kind, then you will have to lie about them as well. Those definitions can be shortened using a generic class:

if TYPE_CHECKING:
    class _PseudoDescriptor[RuntimeType]:
       def __get__(self, instance: Any, owner: type[Any] | None, /) -> RuntimeType: ...

    class IntParameter(_PseudoDescriptor[int]): ...
    class FloatParameter(_PseudoDescriptor[float]): ...

Another way is to make Task a dataclass_transform()-er:

if TYPE_CHECKING:
    @dataclass_transform(
        # Put other classes here
        field_specifiers = (luigi.IntParameter, luigi.FloatParameter)
    )
    class Task: ...
else:
    from luigi import Task

class MyTask(Task):
    int_attr: int = luigi.IntParameter()
    float_attr: float = luigi.FloatParameter()

t = MyTask(int_attr = 42, float_attr = 3.14)
reveal_type(t.int_attr)    # int
reveal_type(t.float_attr)  # float

Simpler alternatives include:

class MyTask(luigi.Task):
    a: int = luigi.IntParameter()  # type: ignore`
    b = cast(int, luigi.IntParameter())

t = MyTask(a = 2, b = 3)
reveal_type(t.a)  # int
reveal_type(t.b)  # int

Personally, I would say the second way is the most elegant overall, but I don't know luigi, so you'll have to decide for yourself which one to use.

Always remember that you can stop using static typing altogether. These workarounds might come back to haunt you in the future. Consider carefully if they are worth the potential maintenance burden.

Upvotes: 1

Related Questions