pythonpython-typingmypyluigi

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.


Solution

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