stubber: Merge overloads, fix adding unused overloads.

Signed-off-by: Jos Verlinde <[email protected]>

src/stubber/codemod/merge_docstub.py

@@ -209,19 +209,22 @@ def add_missed_overloads(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mo
             # insert each overload just after a function with the same name
             # reversed to keep insertions in same order as in the docstub
             for overload in reversed(missing_overloads):
+                matched = False
                 for i, node in enumerate(updated_body):
                     if (
                         isinstance(node, cst.FunctionDef)
                         and node.name.value == overload.name.value
                     ):
+                        matched = True
                         break
-                updated_body.insert(i + 1, overload)
+                if matched:
+                    updated_body.insert(i + 1, overload)
 
             if isinstance(updated_node, cst.Module):
                 updated_node = updated_node.with_changes(body=tuple(updated_body))
             elif isinstance(updated_node, cst.ClassDef):
                 b1 = updated_node.body.with_changes(body=tuple(updated_body))
-                updated_node = updated_node.with_changes( body = b1)
+                updated_node = updated_node.with_changes(body=b1)
 
                 # cst.IndentedBlock(body=tuple(updated_body)))  # type: ignore
         return updated_node

tests/codemods/codemod_test_cases/func_overload_101/output.pyi

@@ -0,0 +1,12 @@
+# fmt: off
+"""
+Overloaded functions from python docs
+"""
+from typing import overload
+
+@overload
+def process(response: None) -> None: ...
+@overload
+def process(response: int) -> tuple[int, str]: ...
+@overload
+def process(response: bytes) -> str: ...

tests/codemods/codemod_test_cases/func_overload_add_2/output.pyi

@@ -0,0 +1,22 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+@overload
+def foo(value: int) -> None:
+    """
+    Set foo value
+    First overload
+    """
+    ...
+
+
+@overload
+def foo(value: None) -> str:
+    """
+    Get foo value
+    Second overload
+    """
+    ...

tests/codemods/codemod_test_cases/func_overload_add_3/output.pyi

@@ -0,0 +1,31 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+@overload
+def foo(value: int) -> None:
+    """
+    Set foo value
+    First overload
+    """
+    ...
+
+
+@overload
+def foo(value: None) -> str:
+    """
+    Get foo value
+    Second overload
+    """
+    ...
+
+
+@overload
+def foo(value: str) -> None:
+    """
+    Get foo string
+    Third overload
+    """
+    ...

tests/codemods/codemod_test_cases/func_overload_add_3_impl/output.pyi

@@ -0,0 +1,34 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+@overload
+def foo(value: int) -> None:
+    """
+    Set foo value
+    First overload
+    """
+    ...
+
+
+@overload
+def foo(value: None) -> str:
+    """
+    Get foo value
+    Second overload
+    """
+    ...
+
+
+@overload
+def foo(value: str) -> None:
+    """
+    Get foo string
+    Third overload
+    """
+    ...
+
+
+def process(): ...

tests/codemods/codemod_test_cases/func_overload_add_3_ordered/output.pyi

@@ -0,0 +1,34 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+@overload
+def foo(value: int) -> None:
+    """
+    Set foo value
+    First overload
+    """
+    ...
+
+
+@overload
+def foo(value: None) -> str:
+    """
+    Get foo value
+    Second overload
+    """
+    ...
+
+
+@overload
+def foo(value: str) -> None:
+    """
+    Get foo string
+    Third overload
+    """
+    ...
+
+
+def bar(): ...

tests/codemods/codemod_test_cases/func_overload_add_no_match/before.pyi

@@ -0,0 +1,6 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+
+def bar(): ...

tests/codemods/codemod_test_cases/func_overload_add_no_match/doc_stub.pyi

@@ -0,0 +1,23 @@
+"""
+Overloaded functions
+"""
+
+from typing import overload
+
+
+@overload
+def foo(value: int) -> None:
+    """
+    Set foo value
+    First overload
+    """
+    ...
+
+
+@overload
+def foo(value: None) -> str:
+    """
+    Get foo value
+    Second overload
+    """
+    ...

tests/codemods/codemod_test_cases/func_overload_add_no_match/expected.pyi

@@ -0,0 +1,7 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+def bar(): ...

tests/codemods/codemod_test_cases/func_overload_add_no_match/output.pyi

@@ -0,0 +1,7 @@
+# fmt: off
+"""
+Overloaded functions
+"""
+from typing import overload
+
+def bar(): ...

tests/codemods/codemod_test_cases/meth_overload_add_2/output.pyi

@@ -0,0 +1,22 @@
+# fmt: off
+"""
+Overloaded methods
+"""
+from typing import overload
+
+class Parrot:
+    @overload
+    def speak(number: int):
+        """
+        Speak a number
+        First overload
+        """
+        ...
+
+    @overload
+    def speak(words: str):
+        """
+        Speak a word
+        Second overload
+        """
+        ...

tests/codemods/codemod_test_cases/meth_overload_add_no_match/before.pyi

@@ -0,0 +1,11 @@
+# fmt: off
+"""
+Overloaded methods
+"""
+
+class Parrot:
+    def talk(): ...
+
+
+class Dog:
+    def speak(): ...

tests/codemods/codemod_test_cases/meth_overload_add_no_match/doc_stub.pyi

@@ -0,0 +1,24 @@
+"""
+Overloaded methods
+"""
+
+from typing import overload
+
+
+class Parrot:
+
+    @overload
+    def speak(number: int):
+        """
+        Speak a number
+        First overload
+        """
+        ...
+
+    @overload
+    def speak(words: str):
+        """
+        Speak a word
+        Second overload
+        """
+        ...

tests/codemods/codemod_test_cases/meth_overload_add_no_match/expected.pyi

@@ -0,0 +1,12 @@
+# fmt: off
+"""
+Overloaded methods
+"""
+from typing import overload
+
+class Parrot:
+    def talk(): ...
+
+
+class Dog:
+    def speak(): ...

tests/codemods/codemod_test_cases/meth_overload_add_no_match/output.pyi

@@ -0,0 +1,12 @@
+# fmt: off
+"""
+Overloaded methods
+"""
+from typing import overload
+
+class Parrot:
+    def talk(): ...
+
+
+class Dog:
+    def speak(): ...