|
3 | 3 | from dataclasses import dataclass |
4 | 4 | from typing import TYPE_CHECKING |
5 | 5 | from typing import Any |
| 6 | +from typing import cast |
6 | 7 |
|
7 | 8 | if TYPE_CHECKING: |
8 | 9 | from collections.abc import Callable |
| 10 | + from typing import Protocol |
9 | 11 |
|
10 | 12 | try: |
11 | 13 | from coiled.function import Function |
12 | 14 | except ImportError: |
13 | 15 |
|
14 | 16 | @dataclass |
15 | 17 | class Function: |
16 | | - cluster_kwargs: dict[str, Any] |
17 | | - environ: dict[str, Any] |
| 18 | + _cluster_kwargs: dict[str, Any] |
| 19 | + _environ: dict[str, Any] | None |
| 20 | + _local: bool |
| 21 | + _name: str |
18 | 22 | function: Callable[..., Any] | None |
19 | | - keepalive: str | None |
| 23 | + keepalive: Any | None |
20 | 24 |
|
21 | 25 |
|
22 | 26 | __all__ = ["Function"] |
23 | 27 |
|
24 | 28 |
|
| 29 | +if TYPE_CHECKING: |
| 30 | + |
| 31 | + class _CoiledFunctionPrivateAttrs(Protocol): |
| 32 | + _cluster_kwargs: dict[str, Any] |
| 33 | + _environ: dict[str, Any] | None |
| 34 | + _local: bool |
| 35 | + _name: str |
| 36 | + keepalive: Any | None |
| 37 | + |
| 38 | + |
| 39 | +def _as_coiled_private_attrs(func: Function) -> _CoiledFunctionPrivateAttrs: |
| 40 | + """Cast to the private-attribute layout used by coiled's Function class.""" |
| 41 | + return cast("_CoiledFunctionPrivateAttrs", func) |
| 42 | + |
| 43 | + |
25 | 44 | def extract_coiled_function_kwargs(func: Function) -> dict[str, Any]: |
26 | 45 | """Extract the kwargs for a coiled function.""" |
| 46 | + coiled_function = _as_coiled_private_attrs(func) |
27 | 47 | return { |
28 | | - "cluster_kwargs": func._cluster_kwargs, # ty: ignore[possibly-missing-attribute] |
29 | | - "keepalive": func.keepalive, |
30 | | - "environ": func._environ, # ty: ignore[possibly-missing-attribute] |
31 | | - "local": func._local, # ty: ignore[possibly-missing-attribute] |
32 | | - "name": func._name, # ty: ignore[possibly-missing-attribute] |
| 48 | + "cluster_kwargs": coiled_function._cluster_kwargs, |
| 49 | + "keepalive": coiled_function.keepalive, |
| 50 | + "environ": coiled_function._environ, |
| 51 | + "local": coiled_function._local, |
| 52 | + "name": coiled_function._name, |
33 | 53 | } |
0 commit comments