Fix bug in task validation to handle variadic args (#98)

This also handles the case of multiple positional only args - since
we only support a single one which is the input dataframe - and where
a keyword argument overlaps with this first argument.

Author: amitschang

dplutils/pipeline/task.py

@@ -1,6 +1,6 @@
+import inspect
 from copy import copy
 from dataclasses import dataclass, field, replace
-from inspect import _empty, signature
 from typing import Any, Callable
 
 import pandas as pd
@@ -77,17 +77,28 @@ def validate(self, context: dict):
         """
         all_kwargs = self.resolve_kwargs(context)
         # we expect a dataframe as the first argument, so skip validation for that
-        params = list(signature(self.func).parameters.items())[1:]
+        params = list(inspect.signature(self.func).parameters.items())[1:]
+        # Because the signature and params therein do not indicate varadics, we have to
+        # consult getfullargspec for those. Similarly, since fullargspec doesn't indicate
+        # positional only, we utilize both.
+        argspec = inspect.getfullargspec(self.func)
+        if argspec.args[0] in all_kwargs:
+            raise ValueError("first position argument reserved for input dataframe but found in kwargs")
         for key, param in params:
-            if param.default == _empty:
+            if param.name in [argspec.varargs, argspec.varkw]:
+                continue
+            if param.kind == inspect.Parameter.POSITIONAL_ONLY:
+                raise ValueError(f"only one positional only argument supported, found also {param.name}")
+            if param.default == inspect._empty:
                 if key not in all_kwargs:
                     msg = f"missing required argument {key} for task {self.name}"
                     if key in self.context_kwargs:
                         msg = f"{msg} - expected from context {self.context_kwargs[key]}"
                     raise ValueError(msg)
-        extra = set(all_kwargs.keys()) - {k for k, v in params}
-        if len(extra) > 0:
-            raise ValueError(f"unkown arguments {extra} for task {self.name}")
+        if not argspec.varkw:
+            extra = set(all_kwargs.keys()) - {k for k, v in params}
+            if len(extra) > 0:
+                raise ValueError(f"unkown arguments {extra} for task {self.name}")
 
     def __hash__(self):
         return hash(self.name)

tests/conftest.py

@@ -93,6 +93,22 @@ def func(dataframe, required, optional=1):
     return PipelineTask("name", func)
 
 
+@pytest.fixture
+def generic_task_with_starkwargs():
+    def func(dataframe, **kwargs):
+        pass
+
+    return PipelineTask("name", func)
+
+
+@pytest.fixture
+def generic_task_with_position_only():
+    def func(dataframe, position_only, /):
+        pass
+
+    return PipelineTask("name", func)
+
+
 @pytest.fixture
 def dummy_executor(dummy_steps):
     return DummyExecutor(graph=dummy_steps)

tests/pipeline/test_pipeline_task.py

@@ -40,5 +40,19 @@ def test_pipeline_task_check_kwargs_from_context(generic_task_with_required):
 
 
 def test_pipeline_task_check_kwargs_from_context_raises_with_missing(generic_task_with_required):
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="expected from context"):
         generic_task_with_required(context_kwargs={"required": "ctx_required"}).validate({})
+
+
+def test_pipeline_task_with_starkwargs_argument_allows_any(generic_task_with_starkwargs):
+    generic_task_with_starkwargs(kwargs={"some_kwarg": 1}).validate({})
+
+
+def test_pipeline_task_with_first_arg_as_kwarg_raises(generic_task_with_starkwargs):
+    with pytest.raises(ValueError, match="first position argument reserved"):
+        generic_task_with_starkwargs(kwargs={"dataframe": 1}).validate({})
+
+
+def test_pipeline_task_with_multiple_position_only_raises(generic_task_with_position_only):
+    with pytest.raises(ValueError, match="only one positional only"):
+        generic_task_with_position_only().validate({})