Support `ctx: UnionType` annots for `@tractor.context` eps

hilevel_serman
Tyler Goodlet 2024-12-11 22:22:26 -05:00
parent 170e198683
commit dcb6706489
1 changed files with 11 additions and 1 deletions

View File

@ -47,6 +47,9 @@ from functools import partial
import inspect import inspect
from pprint import pformat from pprint import pformat
import textwrap import textwrap
from types import (
UnionType,
)
from typing import ( from typing import (
Any, Any,
AsyncGenerator, AsyncGenerator,
@ -2544,7 +2547,14 @@ def context(
name: str name: str
param: Type param: Type
for name, param in annots.items(): for name, param in annots.items():
if param is Context: if (
param is Context
or (
isinstance(param, UnionType)
and
Context in param.__args__
)
):
ctx_var_name: str = name ctx_var_name: str = name
break break
else: else: