Fixed parameter handling bugs for async iterate_over_data_types. output_dataset param is now keyword-only

Author: sveinugu

src/omnipy/compute/mixins/iterate.py

@@ -178,7 +178,8 @@ def _generate_new_signature_for_iteration(self, job_func: Callable) -> None:
             else:
                 output_dataset_param = Parameter(
                     name=self._output_dataset_param,
-                    kind=Parameter.POSITIONAL_OR_KEYWORD,
+                    kind=Parameter.KEYWORD_ONLY,
+                    default=None,
                     annotation=output_dataset_cls,
                 )
                 rest_params = rest_params + [output_dataset_param]
@@ -200,19 +201,28 @@ def _extract_output_dataset(
         self_as_signature_func_job_base_mixin = cast(SignatureFuncJobBaseMixin, self)
 
         if self._output_dataset_param:
+            if self._output_dataset_param_in_func and self._output_dataset_param not in kwargs:
+                kwargs = kwargs.copy()
+                kwargs[self._output_dataset_param] = cast(
+                    Dataset, self_as_signature_func_job_base_mixin.return_type())
+
             bound_args = self_as_signature_func_job_base_mixin.get_bound_args(
                 dataset, *args, **kwargs)
-            output_dataset: Dataset = bound_args.arguments[self._output_dataset_param]
-
-            if self._output_dataset_param_in_func:
-                return_args = bound_args.args[1:]
-            else:
-                return_args = bound_args.args[1:-1]
+            output_dataset: Dataset | None = bound_args.arguments[self._output_dataset_param]
 
-            return output_dataset, return_args, bound_args.kwargs
+            return_args = bound_args.args[1:]
+            return_kwargs = bound_args.kwargs
+            if not self._output_dataset_param_in_func:
+                return_kwargs.pop(self._output_dataset_param)
         else:
+            output_dataset = None
+            return_args = args
+            return_kwargs = kwargs
+
+        if output_dataset is None:
             output_dataset = cast(Dataset, self_as_signature_func_job_base_mixin.return_type())
-            return output_dataset, args, kwargs
+
+        return output_dataset, return_args, return_kwargs
 
     def _prepare_data_arg(self, data_file):
         return data_file if is_model_subclass(self._input_dataset_type) else data_file.contents

tests/compute/cases/iterate_tasks.py

@@ -12,6 +12,7 @@
                             async_single_int_plus_future_return_alphanum_string_func,
                             async_single_int_plus_int_return_str_func,
                             async_single_int_plus_int_return_str_model_with_output_str_dataset_func,
+                            single_int_model_plus_default_int_pair_return_str_model_func,
                             single_int_model_plus_int_return_str_model_func,
                             single_int_plus_int_return_str_func,
                             single_int_plus_int_return_str_model_with_output_str_dataset_func,
@@ -89,6 +90,28 @@ def case_sync_single_int_model_plus_int_return_str_model_func() -> IterateDataFi
     )
 
 
+@pc.case(
+    id='sync_single_int_model_plus_default_int_pair_return_str_model_func',
+    tags=[
+        'sync',
+        'function',
+        'iterate',
+        'no_output_dataset',
+        'str_output_dataset',
+        'int_output_dataset'
+    ],
+)
+def case_sync_single_int_model_plus_default_int_pair_return_str_model_func(
+) -> IterateDataFilesCase:
+    return IterateDataFilesCase(
+        task_func=single_int_model_plus_default_int_pair_return_str_model_func,
+        args=(1,),
+        kwargs={'other_number': 1},
+        func_is_async=False,
+        iterate_over_data_files=True,
+    )
+
+
 @pc.case(
     id='sync_single_int_plus_int_return_str_func',
     tags=[

tests/compute/cases/raw/functions.py

@@ -64,7 +64,15 @@ def single_int_model_plus_int_return_str_model_func(
     data_number: Model[int],
     number: int,
 ) -> Model[str]:
-    return str(data_number.contents + number)
+    return str(data_number.contents + number)  # type: ignore[return-value]
+
+
+def single_int_model_plus_default_int_pair_return_str_model_func(
+    data_number: Model[int],
+    number: int = 0,
+    other_number: int = 0,
+) -> Model[str]:
+    return str(data_number.contents + number + other_number)  # type: ignore[return-value]
 
 
 def single_int_plus_int_return_str_func(data_number: int, number: int) -> str:
@@ -76,7 +84,7 @@ def single_int_plus_int_return_str_model_with_output_str_dataset_func(
     number: int,
     output_dataset: Dataset[Model[str]],
 ) -> Model[str]:
-    return str(data_number + number)
+    return str(data_number + number)  # type: ignore[return-value]
 
 
 def single_int_plus_int_return_str_with_output_int_dataset_func(
@@ -92,7 +100,7 @@ async def async_single_int_model_plus_int_return_str_model_func(
     number: int,
 ) -> Model[str]:
     await asyncio.sleep(random() / 10.0)
-    return str(data_number.contents + number)
+    return str(data_number.contents + number)  # type: ignore[return-value]
 
 
 async def async_single_int_plus_int_return_str_func(data_number: int, number: int) -> str:
@@ -114,7 +122,7 @@ async def async_single_int_plus_int_return_str_model_with_output_str_dataset_fun
     output_dataset: Dataset[Model[str]],
 ) -> Model[str]:
     await asyncio.sleep(random() / 10.0)
-    return str(data_number + number)
+    return str(data_number + number)  # type: ignore[return-value]
 
 
 async def async_single_int_plus_future_int_return_str_func(

tests/compute/mixins/test_iterate.py

@@ -131,6 +131,24 @@ async def test_iterate_over_data_files_task(case: IterateDataFilesCase) -> None:
     _assert_str_result(case, returned_dataset)
 
 
+@pc.parametrize_with_cases(
+    'case', cases='..cases.iterate_tasks', has_tag=['iterate', 'str_output_dataset'])
+async def test_iterate_over_data_files_with_default_output_dataset_param_task(
+        case: IterateDataFilesCase) -> None:
+
+    task_template = TaskTemplate(
+        iterate_over_data_files=case.iterate_over_data_files,
+        output_dataset_param='output_dataset')(
+            case.task_func)
+
+    dataset = Dataset[Model[int]](dict(a=3, b=5, c=-2))
+
+    dataset_or_task = _run_task_template(case, task_template, dataset)
+    returned_dataset = await _ensure_dataset_await_if_task(case, dataset_or_task)
+
+    _assert_str_result(case, returned_dataset)
+
+
 @pc.parametrize_with_cases(
     'case', cases='..cases.iterate_tasks', has_tag=['iterate', 'str_output_dataset'])
 async def test_iterate_over_data_files_with_output_dataset_param_task(

tests/compute/mixins/test_mixin_integration.py

@@ -62,7 +62,8 @@ def test_iterate_over_data_files_func_signature_output_dataset_param() -> None:
             'output_dataset':
                 Parameter(
                     'output_dataset',
-                    Parameter.POSITIONAL_OR_KEYWORD,
+                    Parameter.KEYWORD_ONLY,
+                    default=None,
                     annotation=Dataset[Model[str]])
         }
         assert task_obj.return_type is Dataset[Model[str]]
@@ -85,7 +86,10 @@ def test_iterate_over_data_files_func_signature_output_dataset_param_and_cls() -
                 Parameter('number', Parameter.POSITIONAL_OR_KEYWORD, annotation=int),
             'output_dataset':
                 Parameter(
-                    'output_dataset', Parameter.POSITIONAL_OR_KEYWORD, annotation=CustomStrDataset)
+                    'output_dataset',
+                    Parameter.KEYWORD_ONLY,
+                    default=None,
+                    annotation=CustomStrDataset)
         }
         assert task_obj.return_type is CustomStrDataset