From 42dae86eac7a96d53004d05cbd45141a8eaaa5c1 Mon Sep 17 00:00:00 2001
From: Tobias Raabe <raabe@posteo.de>
Date: Sat, 11 Nov 2023 23:25:51 +0100
Subject: [PATCH] Addition to #489.

---
 docs/source/changes.md       |  6 +++---
 src/_pytask/collect_utils.py | 27 +++++++--------------------
 2 files changed, 10 insertions(+), 23 deletions(-)

diff --git a/docs/source/changes.md b/docs/source/changes.md
index 90cfec84..7ac04f26 100644
--- a/docs/source/changes.md
+++ b/docs/source/changes.md
@@ -14,9 +14,9 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
   caused flaky tests.
 - {pull}`486` adds default names to {class}`~pytask.PPathNode`.
 - {pull}`488` raises an error when an invalid value is used in a return annotation.
-- {pull}`489` simplifies parsing products and does not raise an error when a product
-  annotation is used with the argument name `produces`. And, allow `produces` to intake
-  any node.
+- {pull}`489` and {pull}`491` simplifies parsing products and does not raise an error
+  when a product annotation is used with the argument name `produces`. And, allow
+  `produces` to intake any node.
 - {pull}`490` refactors and better tests parsing of dependencies.
 
 ## 0.4.2 - 2023-11-8
diff --git a/src/_pytask/collect_utils.py b/src/_pytask/collect_utils.py
index 6f306e02..054a01e3 100644
--- a/src/_pytask/collect_utils.py
+++ b/src/_pytask/collect_utils.py
@@ -406,10 +406,16 @@ def parse_products_from_task_function(  # noqa: C901
         if set(parameters_with_product_annot) - {"produces"}:
             has_produces_argument = True
 
+    if "return" in parameters_with_node_annot:
+        parameters_with_product_annot.append("return")
+        has_return = True
+
     if parameters_with_product_annot:
         out = {}
         for parameter_name in parameters_with_product_annot:
-            has_annotation = True
+            if parameter_name != "return":
+                has_annotation = True
+
             if (
                 parameter_name not in kwargs
                 and parameter_name not in parameters_with_node_annot
@@ -448,25 +454,6 @@ def parse_products_from_task_function(  # noqa: C901
             )
             out[parameter_name] = collected_products
 
-    if "return" in parameters_with_node_annot:
-        has_return = True
-        collected_products = tree_map_with_path(
-            lambda p, x: _collect_product(
-                session,
-                node_path,
-                task_name,
-                NodeInfo(
-                    arg_name="return",
-                    path=p,
-                    value=x,
-                    task_path=task_path,
-                    task_name=task_name,
-                ),
-            ),
-            parameters_with_node_annot["return"],
-        )
-        out = {"return": collected_products}
-
     task_produces = obj.pytask_meta.produces if hasattr(obj, "pytask_meta") else None
     if task_produces:
         has_task_decorator = True