feat: collect Pin.__call__ from the docstub of the same version

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

Author: Josverl

src/stubber/codemod/add_method.py

@@ -5,32 +5,38 @@
 """
 
 
+from typing import Optional
 import libcst as cst
 from libcst import matchers as m
 
-# todo: extract the call method from the docstubs
-FNDEF_PIN_CALL = '''
-def __call__(self, x: Optional[Any] = None) -> Any:
-    """
-    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.
-    """
-    ...
-'''
 
-# print(call_meth)
+# there is no simple way to re-use the code from multiple classes / methods 
+# but could be added https://stackoverflow.com/questions/17522706/how-to-pass-an-instance-variable-to-a-decorator-inside-class-definition?noredirect=1&lq=1
+# so for now, just copy the code, or use a module scoped variable - but that is not thread safe
+
+class CallFinder(m.MatcherDecoratableTransformer):
+    """Find the Pin.__call__ method and extract it from a (machine) module."""
+    class_name: str = "Pin"  # class name
+    method_name: str = "__call__"  # method name
+    def __init__(self):
+        super().__init__()
+        self.call_meth: Optional[cst.FunctionDef] = None
+
+    @m.call_if_inside(m.ClassDef(name=m.Name(class_name)))
+    @m.visit(m.FunctionDef(name=m.Name(method_name)))
+    def detect_call(self, node: cst.FunctionDef) -> None:
+        """find the  __call__ method and store it."""
+        self.call_meth = node
 
 
 class CallAdder(m.MatcherDecoratableTransformer):
     """Add a __call__ method to a class if it is missing."""
     class_name = "Pin"  # class name 
     has_call = 0 # number of __call__ methods found
 
-    def __init__(self) -> None:
+    def __init__(self, call_meth :cst.FunctionDef) -> None:
         super().__init__()
-        # parse the (default) call method once
-        self.call_meth = cst.parse_statement(FNDEF_PIN_CALL)
+        self.call_meth  = call_meth
 
 
     @m.call_if_inside(m.ClassDef(name=m.Name(class_name)))

src/stubber/publish/merge_docstubs.py

@@ -100,7 +100,7 @@ def merge_all_docstubs(
         log.info(f"Merge docstubs for {merged_path.name} {candidate['version']}")
         result = copy_and_merge_docstubs(board_path, merged_path, doc_path)
         # Add methods from docstubs to the firmware stubs that do not exist in the firmware stubs
-        add_machine_pin_call(merged_path)
+        add_machine_pin_call(merged_path, candidate['version'])
         if result:
             merged += 1
     log.info(f"merged {merged} of {len(candidates)} candidates")

src/stubber/publish/missing_class_methods.py

@@ -5,20 +5,44 @@
 from pathlib import Path
 
 import libcst as cst
+from loguru import logger as log
 
-from stubber.codemod.add_method import CallAdder
+from stubber.codemod.add_method import CallAdder, CallFinder
+from stubber.utils.config import CONFIG
 from stubber.utils.post import run_black
+from stubber.utils.versions import clean_version
 
 
-def add_machine_pin_call(merged_path: Path):
+def add_machine_pin_call(merged_path: Path, version: str):
     """
     Add the __call__ method to the machine.Pin and pyb.Pin class
     in all pyb and machine/umachine stubs
     """
+    # TODO: this should be done in the merge_docstubs.py to avoid needing to run black twice
+    # and to avoid having to parse the file twice
+
+    # first find the __call__ method in the default stubs
+    mod_path = CONFIG.stub_path / f"micropython-{clean_version(version, flat=True)}-docstubs/machine.pyi"
+    if not mod_path.exists():
+        log.error(f"no docstubs found for {version}")
+        return False
+    log.trace(f"Parsing {mod_path} for __call__ method")
+    source = mod_path.read_text()
+    module = cst.parse_module(source)
+
+    call_finder = CallFinder()
+    module.visit(call_finder)
+
+    if call_finder.call_meth is None:
+        log.error("no __call__ method found")
+        return False
+
+    # then use the CallAdder to add the __call__ method to all machine and pyb stubs
     mod_paths = [f for f in merged_path.rglob("*.*") if f.stem in {"machine", "umachine", "pyb"}]
     for mod_path in mod_paths:
         source = mod_path.read_text()
         machine_module = cst.parse_module(source)
-        new_module = machine_module.visit(CallAdder())
+        new_module = machine_module.visit(CallAdder(call_finder.call_meth))
         mod_path.write_text(new_module.code)
         run_black(mod_path)
+    return True