Support overloads in type converters

KotlinIsland · KotlinIsland · commit afb764593825 · 2023-02-06T15:42:36.000+10:00
diff --git a/atest/robot/keywords/type_conversion/overloads.robot b/atest/robot/keywords/type_conversion/overloads.robot
@@ -0,0 +1,10 @@
+*** Settings ***
+Suite Setup       Run Tests    ${EMPTY}    keywords/type_conversion/overloads.robot
+Resource          atest_resource.robot
+
+*** Test Cases ***
+Annotated
+    Check Test Case    ${TESTNAME}
+
+Unannotated
+    Check Test Case    ${TESTNAME}
diff --git a/atest/testdata/keywords/type_conversion/overloads.py b/atest/testdata/keywords/type_conversion/overloads.py
@@ -0,0 +1,27 @@
+from typing import Union
+
+from typing_extensions import overload
+
+
+@overload
+def foo(argument: int, expected: object): ...
+
+
+@overload
+def foo(argument: None, expected: object): ...
+
+
+def foo(argument: Union[int, None], expected: object):
+    assert argument == expected
+
+
+@overload
+def bar(argument: int, expected: object): ...
+
+
+@overload
+def bar(argument: None, expected: object): ...
+
+
+def bar(argument, expected):
+    assert argument == expected
diff --git a/atest/testdata/keywords/type_conversion/overloads.robot b/atest/testdata/keywords/type_conversion/overloads.robot
@@ -0,0 +1,23 @@
+*** Settings ***
+Library                  overloads.py
+Resource                 conversion.resource
+
+*** Test Cases ***
+Annotated
+    foo    1    ${1}
+    foo    None    ${None}
+    TRY
+        foo    a    a
+    EXCEPT     ValueError: Couldn't convert to any overload:\n \ Argument 'argument' got value 'a' that cannot be converted to integer.\n \ Argument 'argument' got value 'a' that cannot be converted to None.
+        No Operation
+    END
+
+
+Unannotated
+    bar    1    ${1}
+    bar    None    ${None}
+    TRY
+        foo    a    a
+    EXCEPT     ValueError: Couldn't convert to any overload:\n \ Argument 'argument' got value 'a' that cannot be converted to integer.\n \ Argument 'argument' got value 'a' that cannot be converted to None.
+        No Operation
+    END
diff --git a/src/robot/running/arguments/argumentparser.py b/src/robot/running/arguments/argumentparser.py
@@ -21,6 +21,13 @@
 from robot.variables import is_assign, is_scalar_assign
 
 from .argumentspec import ArgumentSpec
+try:
+    from typing import get_overloads
+except ImportError:
+    try:
+        from typing_extensions import get_overloads
+    except ImportError:
+        get_overloads = None
 
 
 class _ArgumentParser:
@@ -41,8 +48,16 @@ def _report_error(self, error):
 
 class PythonArgumentParser(_ArgumentParser):
 
-    def parse(self, handler, name=None):
-        spec = ArgumentSpec(name, self._type)
+    def parse(self, handler, name=None, parse_overloads=True):
+        try:
+            overloads = get_overloads(handler)
+        except Exception:  # < 3.10 doesn't work with `functools.partial` objects
+            overloads = []
+        spec_overloads = (
+            tuple(self.parse(overload, name, False) for overload in overloads)
+            if get_overloads and parse_overloads else ()
+        )
+        spec = ArgumentSpec(name, self._type, overloads=spec_overloads)
         self._set_args(spec, handler)
         self._set_types(spec, handler)
         return spec
diff --git a/src/robot/running/arguments/argumentspec.py b/src/robot/running/arguments/argumentspec.py
@@ -28,7 +28,7 @@ class ArgumentSpec:
 
     def __init__(self, name=None, type='Keyword', positional_only=None,
                  positional_or_named=None, var_positional=None, named_only=None,
-                 var_named=None, defaults=None, types=None):
+                 var_named=None, defaults=None, types=None, overloads=()):
         self.name = name
         self.type = type
         self.positional_only = positional_only or []
@@ -38,6 +38,7 @@ def __init__(self, name=None, type='Keyword', positional_only=None,
         self.var_named = var_named
         self.defaults = defaults or {}
         self.types = types
+        self.overloads = overloads
 
     @setter
     def types(self, types):
diff --git a/src/robot/running/handlers.py b/src/robot/running/handlers.py
@@ -83,8 +83,19 @@ def _get_initial_handler(self, library, name, method):
         return None
 
     def resolve_arguments(self, args, variables=None, languages=None):
-        return self.arguments.resolve(args, variables, self.library.converters,
-                                      languages=languages)
+        if not self.arguments.overloads:
+            return self.arguments.resolve(args, variables, self.library.converters,
+                                          languages=languages)
+        errors = []
+        for spec in self.arguments.overloads:
+            try:
+                return spec.resolve(args, variables, self.library.converters,
+                                    languages=languages)
+            except ValueError as e:
+                errors.append(str(e))
+        message = '\n  '.join(errors)
+        raise ValueError(f"Couldn't convert to any overload:\n  {message}")
+
 
     @property
     def doc(self):
diff --git a/utest/requirements.txt b/utest/requirements.txt
@@ -1,4 +1,4 @@
 # External Python modules required by unit tests.
 docutils >= 0.10
 jsonschema
-typing_extensions; python_version <= '3.8'
+typing_extensions  # needed for backported overloads
