core: ensure function return type can be class of string

Author: 2maz

src/damast/core/decorators.py

@@ -17,6 +17,26 @@
 from .transformations import PipelineElement
 
 
+def _require_annotated_dataframe_return(func, *, decorator_name: str) -> None:
+    """
+    Ensure a decorated function is annotated as returning an :class:`AnnotatedDataFrame`.
+
+    The annotation is accepted both as the class itself and as the string ``'AnnotatedDataFrame'``,
+    since a module using ``from __future__ import annotations`` (PEP 563) only ever hands over
+    strings - the transformer is perfectly valid in that case.
+
+    :param func: the function being decorated
+    :param decorator_name: name of the calling decorator, for the error message
+    :raise RuntimeError: if the function is not annotated as returning an ``AnnotatedDataFrame``
+    """
+    return_type = inspect.signature(func).return_annotation
+    if return_type not in (AnnotatedDataFrame, "AnnotatedDataFrame"):
+        raise RuntimeError(
+            f"{decorator_name}: decorator requires 'AnnotatedDataFrame' to be returned by"
+            f" function - but was '{return_type}'"
+        )
+
+
 def _get_dataframe(*args, **kwargs) -> AnnotatedDataFrame:
     """
     Extract the dataframe from positional or keyword arguments
@@ -169,11 +189,7 @@ def output(requirements: dict[str, Any], exclusive: bool = False):
     )
 
     def decorator(func):
-        return_type = inspect.signature(func).return_annotation
-        if return_type != AnnotatedDataFrame and return_type != 'AnnotatedDataFrame':
-            raise RuntimeError(
-                f"output: decorator requires 'AnnotatedDataFrame' to be returned by function - but was '{return_type}'"
-            )
+        _require_annotated_dataframe_return(func, decorator_name="output")
 
         setattr(func, DECORATED_OUTPUT_SPECS, required_output_specs)
         setattr(func, DECORATED_OUTPUT_EXCLUSIVE, exclusive)
@@ -262,11 +278,7 @@ def artifacts(requirements: dict[str, Any]):
     required_artifact_specs = ArtifactSpecification(requirements=requirements)
 
     def decorator(func):
-        return_type = inspect.signature(func).return_annotation
-        if return_type != AnnotatedDataFrame:
-            raise RuntimeError(
-                "artifacts: decorator requires 'AnnotatedDataFrame' to be returned by function"
-            )
+        _require_annotated_dataframe_return(func, decorator_name="artifacts")
 
         setattr(func, DECORATED_ARTIFACT_SPECS, required_artifact_specs)
 

tests/damast/core/test_dataprocessing.py

@@ -702,6 +702,46 @@ def transform(self, df: AnnotatedDataFrame) -> AnnotatedDataFrame:
     p.transform(df=adf)
 
 
+def test_decorators_accept_string_return_annotation(tmp_path):
+    """
+    A module using 'from __future__ import annotations' (PEP 563) hands the decorators the
+    return annotation as a string - the transformer is valid all the same.
+    """
+    df = polars.LazyFrame([["10000000", 0]], ["mmsi", "status"], orient="row")
+    adf = AnnotatedDataFrame(df, MetaData(columns=[DataSpecification(name="mmsi"),
+                                                   DataSpecification(name="status", unit=units.deg)]))
+
+    class TransformX(PipelineElement):
+
+        @damast.describe("Generic transform of x")
+        @damast.input({"x": {"unit": units.deg}})
+        @damast.output({})
+        @damast.artifacts({"file": "string-annotation.damast"})
+        def transform(self, df: AnnotatedDataFrame) -> "AnnotatedDataFrame":
+            (self.parent_pipeline.base_dir / "string-annotation.damast").write_text("test")
+            return df
+
+    pipeline = DataProcessingPipeline(name="string-annotation-test", base_dir=tmp_path) \
+        .add(name="status-transform", transformer=TransformX(), name_mappings={"x": "status"})
+    pipeline.transform(df=adf)
+
+    assert (tmp_path / "string-annotation.damast").exists()
+
+
+@pytest.mark.parametrize("decorator", [damast.output({}), damast.artifacts({"file": "test.damast"})])
+def test_decorators_reject_other_return_annotation(decorator):
+    """Any other return annotation - string or not - remains an error."""
+    with pytest.raises(RuntimeError, match="AnnotatedDataFrame"):
+        @decorator
+        def transform(self, df: AnnotatedDataFrame) -> polars.LazyFrame:
+            return df
+
+    with pytest.raises(RuntimeError, match="AnnotatedDataFrame"):
+        @decorator
+        def transform_by_name(self, df: AnnotatedDataFrame) -> "polars.LazyFrame":
+            return df
+
+
 def test_save(tmp_path):
     pipeline = DataProcessingPipeline(name="abc",
                                       base_dir=tmp_path) \