From 6e4909ece5b7dec974b7a1435178034b1dd23f0d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mose=20M=C3=BCller?= <mosemueller@gmail.com>
Date: Thu, 10 Aug 2023 10:48:20 +0200
Subject: [PATCH] feat: adding support for union type hint serialization

---
 src/pydase/data_service/data_service.py | 30 +++++++++++++++++++------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/pydase/data_service/data_service.py b/src/pydase/data_service/data_service.py
index f7eef23..7c7d60c 100644
--- a/src/pydase/data_service/data_service.py
+++ b/src/pydase/data_service/data_service.py
@@ -3,7 +3,7 @@ import inspect
 import json
 import os
 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
 from loguru import logger
@@ -265,12 +265,28 @@ class DataService(rpyc.Service, AbstractDataService):
                 }
             elif inspect.isfunction(value) or inspect.ismethod(value):
                 sig = inspect.signature(value)
-                parameters = {
-                    k: v.annotation.__name__
-                    if v.annotation is not inspect._empty
-                    else None
-                    for k, v in sig.parameters.items()
-                }
+                parameters: dict[str, Optional[str]] = {}
+                for k, v in sig.parameters.items():
+                    annotation = v.annotation
+                    if annotation is not inspect._empty:
+                        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
                 if (
                     key in self._task_manager.tasks