Refactor coiled kwargs extraction adapter

Author: tobiasraabe

src/_pytask/coiled_utils.py

@@ -3,49 +3,51 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import cast
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+    from typing import Protocol
 
 try:
     from coiled.function import Function
 except ImportError:
 
     @dataclass
     class Function:
-        cluster_kwargs: dict[str, Any]
-        environ: dict[str, Any]
+        _cluster_kwargs: dict[str, Any]
+        _environ: dict[str, Any] | None
+        _local: bool
+        _name: str
         function: Callable[..., Any] | None
-        keepalive: str | None
+        keepalive: Any | None
 
 
 __all__ = ["Function"]
 
 
-_MISSING = object()
+if TYPE_CHECKING:
+
+    class _CoiledFunctionPrivateAttrs(Protocol):
+        _cluster_kwargs: dict[str, Any]
+        _environ: dict[str, Any] | None
+        _local: bool
+        _name: str
+        keepalive: Any | None
 
 
-def _get_coiled_attribute(func: Function, *names: str, default: Any = _MISSING) -> Any:
-    """Get an attribute from coiled function objects with private/public fallbacks."""
-    for name in names:
-        value = getattr(func, name, _MISSING)
-        if value is not _MISSING:
-            return value
-    if default is not _MISSING:
-        return default
-    names_as_text = ", ".join(repr(name) for name in names)
-    msg = f"Cannot find coiled attribute(s) {names_as_text} on {func!r}."
-    raise AttributeError(msg)
+def _as_coiled_private_attrs(func: Function) -> _CoiledFunctionPrivateAttrs:
+    """Cast to the private-attribute layout used by coiled's Function class."""
+    return cast("_CoiledFunctionPrivateAttrs", func)
 
 
 def extract_coiled_function_kwargs(func: Function) -> dict[str, Any]:
     """Extract the kwargs for a coiled function."""
+    coiled_function = _as_coiled_private_attrs(func)
     return {
-        "cluster_kwargs": _get_coiled_attribute(
-            func, "_cluster_kwargs", "cluster_kwargs"
-        ),
-        "keepalive": func.keepalive,
-        "environ": _get_coiled_attribute(func, "_environ", "environ"),
-        "local": _get_coiled_attribute(func, "_local", "local", default=None),
-        "name": _get_coiled_attribute(func, "_name", "name", default=None),
+        "cluster_kwargs": coiled_function._cluster_kwargs,
+        "keepalive": coiled_function.keepalive,
+        "environ": coiled_function._environ,
+        "local": coiled_function._local,
+        "name": coiled_function._name,
     }