Repository: PrefectHQ/prefect

Issue:
Investigate a splatted version of `defaults_from_attrs`
Re: https://github.com/PrefectHQ/prefect/pull/673#discussion_r258760363

Currently, `defaults_from_attrs` accepts a _list_ of strings for which kwargs relate to attrs.  This is functional, but isn't pretty.  Ideally, `defaults_from_attrs` would accept a splatted list, as in the linked comment. 

A naive implementation of this, however, gets caught in the `SignatureValidator` for Tasks, which doesn't allow the `*args` pattern, so this will require a little meta-programming work.


Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/database/sqlite.py b/src/prefect/tasks/database/sqlite.py
--- a/src/prefect/tasks/database/sqlite.py
+++ b/src/prefect/tasks/database/sqlite.py
@@ -27,7 +27,7 @@ def __init__(self, db: str, query: str = None, **kwargs):
         self.query = query
         super().__init__(**kwargs)
 
-    @defaults_from_attrs(["query"])
+    @defaults_from_attrs("query")
     def run(self, query: str = None):
         """
         Args:
@@ -62,7 +62,7 @@ def __init__(self, db: str, script: str = None, **kwargs):
         self.script = script
         super().__init__(**kwargs)
 
-    @defaults_from_attrs(["script"])
+    @defaults_from_attrs("script")
     def run(self, script: str = None):
         """
         Args:
diff --git a/src/prefect/tasks/google/storage.py b/src/prefect/tasks/google/storage.py
--- a/src/prefect/tasks/google/storage.py
+++ b/src/prefect/tasks/google/storage.py
@@ -107,7 +107,7 @@ def __init__(
         )
 
     @defaults_from_attrs(
-        ["blob", "bucket", "project", "credentials_secret", "encryption_key_secret"]
+        "blob", "bucket", "project", "credentials_secret", "encryption_key_secret"
     )
     def run(
         self,
@@ -200,14 +200,12 @@ def __init__(
         )
 
     @defaults_from_attrs(
-        [
-            "bucket",
-            "blob",
-            "project",
-            "create_bucket",
-            "credentials_secret",
-            "encryption_key_secret",
-        ]
+        "bucket",
+        "blob",
+        "project",
+        "create_bucket",
+        "credentials_secret",
+        "encryption_key_secret",
     )
     def run(
         self,
diff --git a/src/prefect/utilities/tasks.py b/src/prefect/utilities/tasks.py
--- a/src/prefect/utilities/tasks.py
+++ b/src/prefect/utilities/tasks.py
@@ -198,8 +198,7 @@ def __init__(self, task: "prefect.Task"):
         self.task = as_task(task)
 
 
-@curry
-def defaults_from_attrs(attr_args: Iterable[str], run_method: Callable) -> Callable:
+def defaults_from_attrs(*attr_args: str) -> Callable:
     """
     Helper decorator for dealing with Task classes with attributes which serve
     as defaults for `Task.run`.  Specifically, this decorator allows the author of a Task
@@ -209,11 +208,10 @@ def defaults_from_attrs(attr_args: Iterable[str], run_method: Callable) -> Calla
     Task is called.
 
     Args:
-        - attr_args (Iterable[str]): an iterable of strings specifying which
+        - *attr_args (str): a splatted list of strings specifying which
             kwargs should fallback to attributes, if not provided at runtime. Note that
             the strings provided here must match keyword arguments in the `run` call signature,
             as well as the names of attributes of this Task.
-        - run_method (Callable): the `Task.run` method to implement the behavior for
 
     Returns:
         - Callable: the decorated / altered `Task.run` method
@@ -225,7 +223,7 @@ def __init__(self, a=None, b=None):
             self.a = a
             self.b = b
 
-        @defaults_from_attrs(['a', 'b'])
+        @defaults_from_attrs('a', 'b')
         def run(self, a=None, b=None):
             return a, b
 
@@ -237,10 +235,13 @@ def run(self, a=None, b=None):
     ```
     """
 
-    @wraps(run_method)
-    def method(self, *args, **kwargs):
-        for attr in attr_args:
-            kwargs.setdefault(attr, getattr(self, attr))
-        return run_method(self, *args, **kwargs)
+    def wrapper(run_method):
+        @wraps(run_method)
+        def method(self, *args, **kwargs):
+            for attr in attr_args:
+                kwargs.setdefault(attr, getattr(self, attr))
+            return run_method(self, *args, **kwargs)
+
+        return method
 
-    return method
+    return wrapper


Test patch:
diff --git a/docs/test_generate_docs.py b/docs/test_generate_docs.py
--- a/docs/test_generate_docs.py
+++ b/docs/test_generate_docs.py
@@ -15,8 +15,9 @@
     get_class_methods,
     get_source,
 )
-from prefect import task
+from prefect import task, Task
 from prefect.engine.state import State
+from prefect.utilities.tasks import defaults_from_attrs
 
 pytest.mark.skipif(sys.version_info < (3, 6))
 pytestmark = pytest.mark.formatting
@@ -91,6 +92,12 @@ def from_nothing(cls, stuff=None):
         pass
 
 
+class MyTask(Task):
+    @defaults_from_attrs("y", "z")
+    def run(self, x, y=None, z=None):
+        return x, y, z
+
+
 @pytest.mark.parametrize(
     "obj,exp",
     [
@@ -106,6 +113,7 @@ def from_nothing(cls, stuff=None):
         (A.from_nothing, "stuff=None"),
         (CustomException, "x"),
         (NamedException, "*args, **kwargs"),
+        (MyTask.run, "x, y=None, z=None"),
     ],
 )
 def test_format_signature(obj, exp):
@@ -123,6 +131,7 @@ def test_format_signature(obj, exp):
         (varargs_no_default, "*args, iso, **kwargs"),
         (A.run, "*args, b=True, **kwargs"),
         (A.y, "*args, b, **kwargs"),
+        (MyTask.run, "x, y=None, z=None"),
     ],
 )
 def test_format_signature_with_curry(obj, exp):
diff --git a/tests/utilities/test_tasks.py b/tests/utilities/test_tasks.py
--- a/tests/utilities/test_tasks.py
+++ b/tests/utilities/test_tasks.py
@@ -1,3 +1,4 @@
+import inspect
 import pytest
 
 from prefect.core import Flow, Task
@@ -203,7 +204,7 @@ def __init__(self, x=None):
                 self.x = x
                 super().__init__()
 
-            @tasks.defaults_from_attrs(["x"])
+            @tasks.defaults_from_attrs("x")
             def run(self, x=None):
                 "Lil doc"
                 return x
@@ -218,7 +219,7 @@ def __init__(self, x=None, y=None):
                 self.y = y
                 super().__init__()
 
-            @tasks.defaults_from_attrs(["x"])
+            @tasks.defaults_from_attrs("x")
             def run(self, x=None, y=None):
                 return x, y
 
@@ -251,6 +252,23 @@ def test_works_with_multiple_args(self, multitask):
         b = multitask(x=1, y=2)
         assert b.run(y=3, x=55) == (55, 3)
 
+    def test_works_with_mutiple_attrs(self):
+        class TestTask(Task):
+            def __init__(self, x=None, y=None):
+                self.x = x
+                self.y = y
+                super().__init__()
+
+            @tasks.defaults_from_attrs("x", "y")
+            def run(self, x=None, y=None):
+                return x, y
+
+        task = TestTask(x=1, y=2)
+        assert task.run() == (1, 2)
+        assert task.run(x=4) == (4, 2)
+        assert task.run(y=99) == (1, 99)
+        assert task.run(x=None, y=None) == (None, None)
+
     def test_raises_if_attr_wasnt_set_at_init(self):
         """
         It would be nice to raise this at creation time, but unfortunately
@@ -258,7 +276,7 @@ def test_raises_if_attr_wasnt_set_at_init(self):
         """
 
         class Forgot(Task):
-            @tasks.defaults_from_attrs(["x"])
+            @tasks.defaults_from_attrs("x")
             def run(self, x=None):
                 return x