feat: adding support for union type hint serialization

This commit is contained in:
Mose Müller 2023-08-10 10:48:20 +02:00
parent 48dd4d58b7
commit 6e4909ece5

View File

@ -3,7 +3,7 @@ import inspect
import json import json
import os import os
from enum import Enum from enum import Enum
from typing import Any, Optional, cast, get_type_hints from typing import Any, Optional, Union, cast, get_args, get_type_hints
import rpyc import rpyc
from loguru import logger from loguru import logger
@ -265,12 +265,28 @@ class DataService(rpyc.Service, AbstractDataService):
} }
elif inspect.isfunction(value) or inspect.ismethod(value): elif inspect.isfunction(value) or inspect.ismethod(value):
sig = inspect.signature(value) sig = inspect.signature(value)
parameters = { parameters: dict[str, Optional[str]] = {}
k: v.annotation.__name__ for k, v in sig.parameters.items():
if v.annotation is not inspect._empty annotation = v.annotation
else None if annotation is not inspect._empty:
for k, v in sig.parameters.items() if (
} hasattr(annotation, "__origin__")
and annotation.__origin__ is Union
):
# Handle union types
union_types = get_args(annotation)
type_names = [
t.__name__ if isinstance(t, type) else str(t)
for t in union_types
]
parameters[k] = " | ".join(type_names)
elif isinstance(annotation, type):
# Handle regular types
parameters[k] = annotation.__name__
else:
parameters[k] = str(annotation)
else:
parameters[k] = None
running_task_info = None running_task_info = None
if ( if (
key in self._task_manager.tasks key in self._task_manager.tasks