stubber: Enhance merge candidate logic and add tests for candidate se…

…lection

Signed-off-by: Jos Verlinde <Jos.Verlinde@microsoft.com>

src/stubber/codemod/enrich.py

@@ -25,6 +25,8 @@
 #########################################################################################
 @dataclass
 class MergeMatch:
+    """A match between a target and source file to merge docstrings and typehints"""
+
     target: Path
     source: Path
     target_pkg: str
@@ -125,9 +127,16 @@ def source_target_candidates(source: Path, target: Path) -> Generator[MergeMatch
         mm = None
         for t in targets:
             # find the best match
+            if t.stem.startswith("u") and t.stem[1:] in U_MODULES:
+                # skip enriching umodule.pyi files
+                log.debug(f"Skip enriching {t.name}, as it is an u-module")
+                continue
             t_pkg = package_from_path(t)
             s_pkg = package_from_path(s)
             is_match, match_len = upackage_equal(s_pkg, t_pkg)
+            if "_mpy_shed" in str(s) or "_mpy_shed" in str(t):
+                log.debug(f"Skip _mpy_shed file {s}")
+                continue
             if is_match and match_len > best_match_len:
                 best_match_len = match_len
                 mm = MergeMatch(t, s, t_pkg, s_pkg, is_match)
@@ -201,6 +210,59 @@ def enrich_file(
         yield diff_code(old_code, current_code, 5, filename=target_path.name)
 
 
+def merge_candidates(
+    source_folder: Path,
+    target_folder: Path,
+) -> List[MergeMatch]:
+    """
+    Generate a list of merge candidates for the source and target folders.
+    Each target is matched with exactly one source file.
+    """
+    candidates = source_target_candidates(source_folder, target_folder)
+
+    # Create a list of candidate matches for the same target
+    target_dict = {}
+    for candidate in candidates:
+        if candidate.target not in target_dict:
+            target_dict[candidate.target] = []
+        target_dict[candidate.target].append(candidate)
+
+    # first get targets with only one candidate
+    candidates = [v[0] for k, v in target_dict.items() if len(v) == 1]
+
+    # then get the best matching from the d
+    multiple_candidates = {k: v for k, v in target_dict.items() if len(v) > 1}
+    for target in multiple_candidates.keys():
+
+        # if simple module --> complex module : select the best matching or first source
+        perfect = next(
+            (
+                match
+                for match in multiple_candidates[target]
+                if match.target_pkg == match.source_pkg
+            ),
+            None,
+        )
+
+        if perfect:
+            candidates.append(perfect)
+        else:
+            close_enough = [
+                match
+                for match in multiple_candidates[target]
+                if match.source_pkg.startswith(f"{match.target_pkg}.")
+            ]
+            if close_enough:
+                candidates.extend(close_enough)
+            # else:
+            #     # take the first one
+            #     candidates.append(multiple_candidates[target][0])
+
+    # sort by target_path , to show diffs
+    candidates = sorted(candidates, key=lambda m: m.target)
+    return candidates
+
+
 def enrich_folder(
     source_folder: Path,
     target_folder: Path,
@@ -231,10 +293,6 @@ def enrich_folder(
 
     # for target in target_files:
     for mm in candidates:
-        if mm.target.stem.startswith("u") and mm.target.stem[1:] in U_MODULES:
-            # skip enriching umodule.pyi files
-            log.debug(f"Skip enriching {mm.target.name}, as it is an u-module")
-            continue
         try:
 
             if diff := list(

src/stubber/codemod/merge_docstub.py

@@ -253,6 +253,10 @@ def add_missed_overloads(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mo
                         # Use the new overload - but with the existing docstring
                         overload = update_def_docstr(overload, docstring_node)
                     updated_body.insert(i + 1, overload)
+                else:
+                    # add to the end of the class
+                    log.trace(f"Add @overload for {overload.name.value} at the end of the class")
+                    updated_body.append(overload)
 
             if isinstance(updated_node, cst.Module):
                 updated_node = updated_node.with_changes(body=tuple(updated_body))

src/stubber/rst/lookup.py

@@ -293,7 +293,7 @@ class Fix:
     ],
     "io": ANY_BUF
     + [
-        "from _mpy_shed import IOBase",
+        "from _mpy_shed import IOBase_mp",
     ],
     "lcd160cr": ANY_BUF + ["from pyb import SPI"],  # uses SPI
     # "machine": ["from network import AbstractNIC"],  # NIC is an abstract class, although not defined or used as such
@@ -652,8 +652,9 @@ class Fix:
     "TextIOWrapper": "IO",  # "TextIOBase, TextIO",  # based on Stdlib
     "FileIO": "IO",  #  "RawIOBase, BinaryIO",  # based on Stdlib
     "StringIO": "IO",  #  "BufferedIOBase, BinaryIO",  # based on Stdlib
+    "IOBase": "IOBase_mp",  # "BufferedIOBase, BinaryIO",  # based on Stdlib
     "BytesIO": "IO",  # "BufferedIOBase, BinaryIO",  # based on Stdlib
-    "BufferedWriter": "IOBase",  # DOC_ME: not in documentation #   "BufferedWriter": "BufferedIOBase",  # based on Stdlib
+    "BufferedWriter": "IOBase_mp",  # DOC_ME: not in documentation #   "BufferedWriter": "BufferedIOBase",  # based on Stdlib
     # uzlib
     # "DecompIO": "IO",  # https://docs.python.org/3/library/typing.html#other-concrete-types
     # -------------------------------------------------------------------------------------

tests/codemods/codemod_test_cases/func_overload_add_call/before.pyi

@@ -33,10 +33,5 @@ class Pin:
     def value(self, *args, **kwargs) -> Incomplete: ...
     def high(self, *args, **kwargs) -> Incomplete: ...
 
-    def __call__(self, x: Optional[Any] = None) -> Incomplete:
-        """
-        Pin objects are callable.  The call method provides a (fast) shortcut to set
-        and get the value of the pin.  It is equivalent to Pin.value([x]).
-        See :meth:`Pin.value` for more details.
-        """
-        ...
+    def foo() -> None: ...
+    # no def __call__ has been defined in the original file

tests/codemods/codemod_test_cases/func_overload_add_call/expected.pyi

@@ -32,6 +32,8 @@ class Pin:
     def init(self, *args, **kwargs) -> Incomplete: ...
     def value(self, *args, **kwargs) -> Incomplete: ...
     def high(self, *args, **kwargs) -> Incomplete: ...
+
+    def foo() -> None: ...
     @overload
     def __call__(self) -> int:
         """
@@ -47,3 +49,4 @@ class Pin:
         and get the value of the pin.  It is equivalent to Pin.value([x]).
         See :meth:`Pin.value` for more details.
         """
+    # no def __call__ has been defined in the original file

tests/codemods/enrich_test.py

@@ -16,12 +16,7 @@
             Path("./tests/data/stub_merge/micropython-v1_18-docstubs/esp32.pyi"),
             Path("./tests/data/stub_merge/micropython-v1_18-esp32/esp32.pyi"),
             True,
-        ),
-        (
-            Path("./tests/data/stub_merge/micropython-v1_18-esp32/builtins.pyi"),
-            Path("./tests/data/stub_merge/micropython-v1_18-esp32/builtins.pyi"),
-            False,
-        ),
+        )
     ],
 )
 def test_enrich_file_with_stub(source_file: Path, target_file: Path, expected: bool):

tests/merge/test_candidates.py

@@ -0,0 +1,76 @@
+from pathlib import Path
+from typing import List, Union
+
+import pytest
+from stubber.codemod.enrich import merge_candidates
+from stubber.publish.candidates import board_candidates
+
+pytestmark = [pytest.mark.stubber]
+
+
+@pytest.mark.parametrize(
+    "family, versions",
+    [
+        ("micropython", "v1.22.0"),
+        ("micropython", "preview"),
+        # Add more test cases here
+    ],
+)
+def test_board_candidates(family: str, versions: Union[str, List[str]]):
+
+    candidates = list(board_candidates(family=family, versions=versions))
+    assert len(candidates) > 0
+
+
+@pytest.mark.parametrize(
+    "id, source, target, count",
+    [
+        (
+            10,
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs",
+            "tests/data/stub_merge/micropython-v1_24_1-rp2-RPI_PICO",
+            55,
+        ),
+        (
+            11,
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs",
+            "tests/data/stub_merge/micropython-v1_24_1-rp2-RPI_PICO/machine.pyi",
+            18,
+        ),
+        (
+            12,
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs",
+            "tests/data/stub_merge/micropython-v1_24_1-rp2-RPI_PICO/umachine.pyi",
+            0,
+        ),
+        (
+            13,
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs",
+            "tests/data/stub_merge/micropython-v1_24_1-rp2-RPI_PICO/micropython.pyi",
+            1,
+        ),
+        (
+            23,
+            "repos/micropython-stubs/micropython-reference",
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs",
+            94,
+        ),
+        (
+            24,
+            "repos/micropython-stubs/micropython-reference",
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs/machine",
+            17,
+        ),
+        (
+            25,
+            "repos/micropython-stubs/micropython-reference",
+            "tests/data/stub_merge/micropython-v1_24_1-docstubs/machine/__init__.pyi",
+            1,
+        ),
+        # Add more test cases here
+    ],
+)
+def test_merge_candidates(id, source, target, count):
+
+    result = merge_candidates(Path(source), Path(target))
+    assert len(result) == count