reworked parse_param_type() function to increase performance, generality and properly handle types: Union[Type1, Type2, ...], __main__.SomeCustomClass

Author: ds-jakub-cierocki

src/dbally/views/exposed_functions.py

@@ -1,26 +1,46 @@
-import re
 from dataclasses import dataclass
+from inspect import isclass
 from typing import _GenericAlias  # type: ignore
 from typing import Optional, Sequence, Type, Union
 
+import typing_extensions as type_ext
+
 from dbally.context.context import BaseCallerContext
 from dbally.similarity import AbstractSimilarityIndex
 
 
-def parse_param_type(param_type: Union[type, _GenericAlias]) -> str:
+def parse_param_type(param_type: Union[type, _GenericAlias, str]) -> str:
     """
     Parses the type of a method parameter and returns a string representation of it.
 
     Args:
-        param_type: type of the parameter
+        param_type: Type of the parameter.
 
     Returns:
-        str: string representation of the type
+        A string representation of the type.
     """
-    if param_type in {int, float, str, bool, list, dict, set, tuple}:
+
+    # TODO consider using hasattr() to ensure correctness of the IF's below
+    if isclass(param_type):
         return param_type.__name__
-    if param_type.__module__ == "typing":
-        return re.sub(r"\btyping\.", "", str(param_type))
+
+    # typing.Literal['aaa', 'bbb'] edge case handler
+    # the args are strings not types thus isclass('aaa') is False
+    # at the same type string has no __module__ property which causes an error
+    if isinstance(param_type, str):
+        return f"'{param_type}'"
+
+    if param_type.__module__ == "typing" or param_type.__module__ == "typing_extensions":
+        type_args = type_ext.get_args(param_type)
+        if type_args:
+            param_name = param_type._name  # pylint: disable=protected-access
+            if param_name is None:
+                # workaround for typing.Literal, because: `typing.Literal['aaa', 'bbb']._name is None`
+                # but at the same time: `type_ext.get_origin(typing.Literal['aaa', 'bbb'])._name == "Literal"`
+                param_name = type_ext.get_origin(param_type)._name  # pylint: disable=protected-access
+
+            args_str_repr = ", ".join(parse_param_type(arg) for arg in type_args)
+            return f"{param_name}[{args_str_repr}]"
 
     return str(param_type)