Repository: PrefectHQ/prefect

Issue:
Keep original docstring from @task decorated function
## Current behavior
When decorating a function with `@task`, the original docstring of the function is lost.

## Proposed behavior
Original docstring is kept. Looking at the wrapper [`prefect.utilities.tasks.task`](https://github.com/PrefectHQ/prefect/blob/c1686091731144bbc0fab367e78b93221fccf28b/src/prefect/utilities/tasks.py#L152-L202), I thinks this should be do-able with [`functools.wraps`](https://docs.python.org/3/library/functools.html#functools.wraps), although I am not exactly sure what change needs to be made.



Discussion:
FYI @dkapitan you can recover the original docstring as follows:

```python
from prefect import task

@task
def my_doc_fun():
    """
    Custom docstring
    """
    pass

my_doc_fun.run.__doc__ # "Custom docstring"
```

Definitely open to elevating this more prominently though (maybe on `my_doc_fun.__doc__` in this case)
I thought this could be achieved by using `functools.update_wrapper`, like:
```python
     task_wrapper = prefect.tasks.core.function.FunctionTask(fn=fn, **task_init_kwargs)
     return update_wrapper(task_wrapper, fn, updated=[])
```
But given `FunctionTask` is a class this doesn't work as expected :-/
@cicdw thanks for clarifying, good to know the original docstring is accessible in the run method.
My use case is that I want to generate documentation using sphinx, so I can deliver that as part of my code to clients to look up what processing steps are done within the ETL. See https://dkapitan.github.io/nimbletl/_build/html/index.html for a first rough setup.

As you can see, the module tasks is empty while it does have FunctionTasks in it.

Any ideas how to fix the (temporary and permanent) are welcome. If anyone can propose a solution I am happy to submit a PR.
I can't promise this will work with Sphinx, but you might try:

```python
from prefect import task

@task
def my_doc_fun():
    """
    Custom docstring
    """
    pass

my_doc_fun.__doc__ = my_doc_run.run.__doc__

## test that this works as far as Python is concerned
import inspect

inspect.getdoc(my_doc_fun) # "Custom docstring"
```

It appears that IPython still considers the original class docstring to be the "true" docstring though, so this might not work.
@cicdw just tried exactly that, no joy ... 
long-running issue in the Sphinx project: https://github.com/sphinx-doc/sphinx/issues/3783. Am trying to figure out how to decorate `FunctionTask` with `functools.wraps`. Surely that must be possible? Beyond my current Python-skills, to looking up SO: https://stackoverflow.com/questions/6394511/python-functools-wraps-equivalent-for-classes
easiest workaround for now:
* define non-decorated function in the library
* declare `task_from_function = task(function)` when writings tasks and flows
The issue here is that prefect task objects aren't functions, they're objects. So for sphinx you'll want to document them as `data` (or `autodata`). Since `automodule` only finds `function` and `class` attributes (not data), you need to document these manually. For example:

```
.. autodata:: nimbletl.tasks.curl_cmd
```

*Note: you can use `autofunction` instead, but that hides that they're task objects that support methods like `map`. Either works though, the important thing is that you have to tell sphinx to document them explicitly, since `automodule` won't find them.*

This doesn't give a satisfying docstring though. Applying this diff to prefect fixes that:

```diff
diff --git a/src/prefect/tasks/core/function.py b/src/prefect/tasks/core/function.py
index ba0fa631d..dfa275bb1 100644
--- a/src/prefect/tasks/core/function.py
+++ b/src/prefect/tasks/core/function.py
@@ -10,9 +10,20 @@ from typing import Any, Callable
 import prefect
 
 
+class _DocProxy(object):
+    def __init__(self, source):
+        self._source = source
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            return self._source
+        else:
+            return obj.run.__doc__
+
+
 class FunctionTask(prefect.Task):
-    """
-    A convenience Task for functionally creating Task instances with
+    __doc__ = _DocProxy(
+    """A convenience Task for functionally creating Task instances with
     arbitrary callable `run` methods.
 
     Args:
@@ -32,7 +43,7 @@ class FunctionTask(prefect.Task):
     with Flow("My Flow") as f:
         result = task(42)
     ```
-    """
+    """)
 
     def __init__(self, fn: Callable, name: str = None, **kwargs: Any):
         if not callable(fn):
```

This sets the docstring of any `@task` wrapped `Task` object to the docstring of the original task, but keeps the docstring for `FunctionTask` around when working with the raw class. I'm not sure if this is desirable though, as `@task` decorated functions are no longer functions (they're `Task` objects), so the docstring isn't 100% accurate. Perhaps we should add a prefix to the function docstring, something like:

```python
"""A prefect.Task wrapping the following function:

{original_docstring}
"""
```

:shrug:

Either way, with the above changes, I get the following for your docs:

<img width="729" alt="Screen Shot 2020-05-18 at 4 12 49 PM" src="https://user-images.githubusercontent.com/2783717/82260211-750dd500-9922-11ea-9883-40c8cfddd42b.png">
I think I'm +0.5 on doing something like the above, just to make interactive development with prefect `@task` wrapped functions more pleasant. For documentation we may want to write a sphinx extension for documenting tasks (e.g. `autotask`) in a nice way. I did this for documenting `traitlet` configuration in `dask-gateway`, and the [required code was fairly short](https://github.com/dask/dask-gateway/blob/master/docs/sphinxext/autodoc_traitlets.py).

Resolving patch:
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -415,6 +415,19 @@ def copy(self, **task_args: Any) -> "Task":
 
         return new
 
+    @property
+    def __signature__(self) -> inspect.Signature:
+        """Dynamically generate the signature, replacing ``*args``/``**kwargs``
+        with parameters from ``run``"""
+        if not hasattr(self, "_cached_signature"):
+            sig = inspect.Signature.from_callable(self.run)
+            parameters = list(sig.parameters.values())
+            parameters.extend(EXTRA_CALL_PARAMETERS)
+            self._cached_signature = inspect.Signature(
+                parameters=parameters, return_annotation="Task"
+            )
+        return self._cached_signature
+
     def __call__(
         self,
         *args: Any,
@@ -1170,3 +1183,12 @@ def serialize(self) -> Dict[str, Any]:
             - dict representing this parameter
         """
         return prefect.serialization.task.ParameterSchema().dump(self)
+
+
+# All keyword-only arguments to Task.__call__, used for dynamically generating
+# Signature objects for Task objects
+EXTRA_CALL_PARAMETERS = [
+    p
+    for p in inspect.Signature.from_callable(Task.__call__).parameters.values()
+    if p.kind == inspect.Parameter.KEYWORD_ONLY
+]
diff --git a/src/prefect/tasks/core/function.py b/src/prefect/tasks/core/function.py
--- a/src/prefect/tasks/core/function.py
+++ b/src/prefect/tasks/core/function.py
@@ -1,8 +1,8 @@
 """
 The tasks in this module can be used to represent arbitrary functions.
 
-In general, users will not instantiate these tasks by hand; they will automatically be
-applied when users apply the `@task` decorator.
+In general, users will not instantiate these tasks by hand; they will
+automatically be applied when users apply the `@task` decorator.
 """
 
 from typing import Any, Callable
@@ -10,9 +10,23 @@
 import prefect
 
 
+class _DocProxy(object):
+    """A descriptor that proxies through the docstring for the wrapped task as
+    the docstring for a `FunctionTask` instance."""
+
+    def __init__(self, cls_doc):
+        self._cls_doc = cls_doc
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            return self._cls_doc
+        else:
+            return getattr(obj.run, "__doc__", None) or self._cls_doc
+
+
 class FunctionTask(prefect.Task):
-    """
-    A convenience Task for functionally creating Task instances with
+    __doc__ = _DocProxy(
+        """A convenience Task for functionally creating Task instances with
     arbitrary callable `run` methods.
 
     Args:
@@ -33,6 +47,7 @@ class FunctionTask(prefect.Task):
         result = task(42)
     ```
     """
+    )
 
     def __init__(self, fn: Callable, name: str = None, **kwargs: Any):
         if not callable(fn):
@@ -46,3 +61,8 @@ def __init__(self, fn: Callable, name: str = None, **kwargs: Any):
         self.run = fn
 
         super().__init__(name=name, **kwargs)
+
+    def __getattr__(self, k):
+        if k == "__wrapped__":
+            return self.run
+        raise AttributeError(f"'FunctionTask' object has no attribute {k}")


Test patch:
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -1,3 +1,4 @@
+import inspect
 import logging
 import uuid
 from datetime import timedelta
@@ -232,6 +233,23 @@ def test_class_instantiation_raises_helpful_warning_for_unsupported_callables(se
         with pytest.raises(ValueError, match="This function can not be inspected"):
             task(zip)
 
+    def test_task_signature_generation(self):
+        class Test(Task):
+            def run(self, x: int, y: bool, z: int = 1):
+                pass
+
+        t = Test()
+
+        sig = inspect.signature(t)
+        # signature is a superset of the `run` method
+        for k, p in inspect.signature(t.run).parameters.items():
+            assert sig.parameters[k] == p
+        # extra kwonly args to __call__ also in sig
+        assert set(sig.parameters).issuperset(
+            {"mapped", "task_args", "upstream_tasks", "flow"}
+        )
+        assert sig.return_annotation == "Task"
+
     def test_create_task_with_and_without_cache_for(self):
         t1 = Task()
         assert t1.cache_validator is never_use
diff --git a/tests/tasks/test_core.py b/tests/tasks/test_core.py
--- a/tests/tasks/test_core.py
+++ b/tests/tasks/test_core.py
@@ -46,6 +46,31 @@ def my_fn():
         f = FunctionTask(fn=my_fn, name="test")
         assert f.name == "test"
 
+    def test_function_task_docstring(self):
+        def my_fn():
+            """An example docstring."""
+            pass
+
+        # Original docstring available on class
+        assert "FunctionTask" in FunctionTask.__doc__
+
+        # Wrapped function is docstring on instance
+        f = FunctionTask(fn=my_fn)
+        assert f.__doc__ == my_fn.__doc__
+
+        # Except when no docstring on wrapped function
+        f = FunctionTask(fn=lambda x: x + 1)
+        assert "FunctionTask" in f.__doc__
+
+    def test_function_task_sets__wrapped__(self):
+        def my_fn():
+            """An example function"""
+            pass
+
+        t = FunctionTask(fn=my_fn)
+        assert t.__wrapped__ == my_fn
+        assert not hasattr(FunctionTask, "__wrapped__")
+
 
 class TestCollections:
     def test_list_returns_a_list(self):