diff --git a/tractor/_context.py b/tractor/_context.py index 4a20c2b..d17f0b1 100644 --- a/tractor/_context.py +++ b/tractor/_context.py @@ -47,6 +47,9 @@ from functools import partial import inspect from pprint import pformat import textwrap +from types import ( + UnionType, +) from typing import ( Any, AsyncGenerator, @@ -2544,7 +2547,14 @@ def context( name: str param: Type 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 break else: