From 02f7c4eac01da5ad86b6ab34398f41d5c988e9e1 Mon Sep 17 00:00:00 2001
From: Alice Purcell <alicederyn@gmail.com>
Date: Wed, 25 Sep 2024 20:17:06 +0100
Subject: [PATCH] Fix type annotations on template decorators

Use `Concatenate` to add the missing `self` parameter to the class methods
annotated with `@_add_type_hints` in `TemplateDecoratorFuncsMixin`. Without
this, mypy raises a `misc` error on usage, complaining that the method "does
not accept self argument".

Signed-off-by: Alice Purcell <alicederyn@gmail.com>
---
 src/hera/workflows/_meta_mixins.py         |  9 ++--
 tests/typehints/test_template_decorator.py | 55 ++++++++++++++++++++++
 2 files changed, 61 insertions(+), 3 deletions(-)
 create mode 100644 tests/typehints/test_template_decorator.py

diff --git a/src/hera/workflows/_meta_mixins.py b/src/hera/workflows/_meta_mixins.py
index 915fbc8a0..319744e6c 100644
--- a/src/hera/workflows/_meta_mixins.py
+++ b/src/hera/workflows/_meta_mixins.py
@@ -13,14 +13,14 @@
 if sys.version_info >= (3, 10):
     from inspect import get_annotations
     from types import NoneType
+    from typing import Concatenate, ParamSpec
 else:
     from hera.shared._inspect import get_annotations
+    from typing_extensions import Concatenate, ParamSpec
 
     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
@@ -474,7 +474,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
diff --git a/tests/typehints/test_template_decorator.py b/tests/typehints/test_template_decorator.py
new file mode 100644
index 000000000..b7dadaf36
--- /dev/null
+++ b/tests/typehints/test_template_decorator.py
@@ -0,0 +1,55 @@
+from pathlib import Path
+from subprocess import run
+from tempfile import TemporaryDirectory
+from textwrap import dedent
+from typing import Iterator, Tuple
+
+import pytest
+
+from hera.workflows import Step, Task
+
+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