materialsproject / jobflow

jobflow is a library for writing computational workflows.
https://materialsproject.github.io/jobflow
Other
90 stars 25 forks source link

Add type hint for @job and fix for run_locally #578

Closed FabiPi3 closed 3 months ago

FabiPi3 commented 3 months ago

Since mypy deprecates the implicit use of None for default values, add the explicit None type hint for store in the run_locally method.

The @job decorator is supposed to transform anything into a Job, I think my change is the minimal effort to take to convince type checkers like mypy or pycharm integrated tools that the result will be a Job. Consider the following examples:

def add(a: int, b: int) -> int:
    """Add func."""
    return a + b

@job
def add_job(a: int, b: int) -> int:
    """Add job."""
    return a + b

@job(large_data="data", graphs="graph")
def compute(a: int, b: int) -> dict[str, int]:
    """Compute job."""
    return {"data": b, "graph": a}

reveal_type(add(1, 2))
reveal_type(add_job(1, 2))
reveal_type(compute(1, 2))

Running mypy over this leads to:

mypy_example.py:15: note: Revealed type is "builtins.int"
mypy_example.py:16: note: Revealed type is "jobflow.core.job.Job"
mypy_example.py:17: note: Revealed type is "jobflow.core.job.Job"

Which is the intended behaviour in my opinion.

utf commented 3 months ago

This is really nice. Thank you.