Skip to content

Commit

Permalink
Fix type annotations on template decorators (#1211)
Browse files Browse the repository at this point in the history
**Pull Request Checklist**
- [X] Fixes #1182
- [X] Tests added
- [ ] Documentation/examples added
- [X] [Good commit messages](https://cbea.ms/git-commit/) and/or PR
title

**Description of PR**
Currently, mypy raises a `misc` error when using a decorator from
`TemplateDecoratorFuncsMixin`, complaining that the method "does not
accept self argument".

This PR uses `Concatenate` to add the missing parameter.

Signed-off-by: Alice Purcell <alicederyn@gmail.com>
  • Loading branch information
alicederyn authored Sep 27, 2024
1 parent 6dea908 commit 005db40
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 3 deletions.
10 changes: 7 additions & 3 deletions src/hera/workflows/_meta_mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,15 @@
if sys.version_info >= (3, 10):
from inspect import get_annotations
from types import NoneType
from typing import Concatenate, ParamSpec
else:
from typing_extensions import Concatenate, ParamSpec

from hera.shared._inspect import get_annotations

NoneType = type(None)


from typing_extensions import ParamSpec

from hera.shared import BaseMixin, global_config
from hera.shared._global_config import _DECORATOR_SYNTAX_FLAG, _flag_enabled
from hera.shared._pydantic import BaseModel, get_fields, root_validator
Expand Down Expand Up @@ -474,7 +475,10 @@ def _add_type_hints(
) -> Callable[
...,
Callable[
PydanticKwargs, # this adds type hints to the underlying *library* function kwargs
Concatenate[
object, # this is the `self` parameter
PydanticKwargs, # this adds type hints to the underlying *library* function kwargs
],
Callable[ # we will return a function that is a decorator
[Callable[FuncIns, FuncR]], # taking underlying *user* function
Callable[FuncIns, FuncR], # and returning it
Expand Down
50 changes: 50 additions & 0 deletions tests/typehints/test_template_decorator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
from pathlib import Path
from subprocess import run
from tempfile import TemporaryDirectory
from textwrap import dedent

COMMON_SETUP = """
from hera.shared import global_config
from hera.workflows import Input, Output, Workflow
w = Workflow()
"""


def run_mypy(python_code: str):
with TemporaryDirectory() as d:
python_file = Path(d) / "example.py"
python_file.write_text(python_code)
mypy_cmd = ["mypy", "--config-file", "tests/typehints/test-mypy.toml", str(python_file)]
result = run(mypy_cmd, check=False, capture_output=True, encoding="utf-8")
if result.returncode != 0:
msg = f"Error calling {' '.join(mypy_cmd)}:\n{result.stderr}{result.stdout}"
raise AssertionError(msg)
return result.stdout.replace(d, "")


def test_script_decoration_no_arguments():
"""Verify a script can be decorated with no extra arguments."""
SCRIPT = """
@w.script()
def simple_script(_: Input) -> Output:
return Output()
reveal_type(simple_script(Input()))
"""
result = run_mypy(COMMON_SETUP + dedent(SCRIPT))
assert 'Revealed type is "hera.workflows.io.v2.Output"' in result


def test_script_decoration_accepts_name_argument():
"""Verify the script decorator can be passed a name."""
SCRIPT = """
@w.script(name = "some_script")
def simple_script(_: Input) -> Output:
return Output()
reveal_type(simple_script(Input()))
"""
result = run_mypy(COMMON_SETUP + dedent(SCRIPT))
assert 'Revealed type is "hera.workflows.io.v2.Output"' in result

0 comments on commit 005db40

Please sign in to comment.