From 5cb26c38e06606b84f515ac7f8f0506ddf057e8b Mon Sep 17 00:00:00 2001
From: Braden Mars <bradenmars@bradenmars.me>
Date: Sun, 5 Mar 2023 05:04:25 -0600
Subject: [PATCH] feat(main): support create stubs backend parameter, utilize
 create stub variant.

Signed-off-by: Braden Mars <bradenmars@bradenmars.me>
---
 micropy/main.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/micropy/main.py b/micropy/main.py
index fa035b47..96a872ea 100644
--- a/micropy/main.py
+++ b/micropy/main.py
@@ -4,15 +4,17 @@
 
 import tempfile
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Literal, Optional, Union
 
 from micropy import data, utils
 from micropy.exceptions import PyDeviceError
-from micropy.lib.stubber import process as stubber
 from micropy.logger import Log
 from micropy.project import Project, modules
 from micropy.pyd import DevicePath, MessageHandlers, ProgressStreamConsumer, PyDevice
+from micropy.pyd.backend_rshell import RShellPyDeviceBackend
+from micropy.pyd.backend_upydevice import UPyDeviceBackend
 from micropy.stubs import RepositoryInfo, StubManager, StubRepository
+from micropy.utils.stub import prepare_create_stubs
 from pydantic import parse_file_as
 
 
@@ -90,7 +92,12 @@ def resolve_project(self, path, verbose=True):
             return proj
         return proj
 
-    def create_stubs(self, port, verbose=False):
+    def create_stubs(
+        self,
+        port,
+        verbose=False,
+        backend: Union[Literal["upydevice"], Literal["rshell"]] = "upydevice",
+    ):
         """Create and add stubs from Pyboard.
 
         Todo:
@@ -105,6 +112,7 @@ def create_stubs(self, port, verbose=False):
         """
         self.log.title(f"Connecting to Pyboard @ $[{port}]")
         pyb_log = Log.add_logger("Pyboard", "bright_white")
+        backend = UPyDeviceBackend if backend == "upydevice" else RShellPyDeviceBackend
 
         def _get_desc(name: str, cfg: dict):
             desc = f"{pyb_log.get_service()} {name}"
@@ -119,15 +127,16 @@ def _get_desc(name: str, cfg: dict):
                 auto_connect=True,
                 stream_consumer=ProgressStreamConsumer(on_description=_get_desc),
                 message_consumer=message_handler,
+                backend=backend,
             )
         except (SystemExit, PyDeviceError):
             self.log.error(f"Failed to connect, are you sure $[{port}] is correct?")
             return None
         self.log.success("Connected!")
-        script = stubber.minify_script(stubber.source_script)
+        create_stubs = prepare_create_stubs()
         self.log.info("Executing stubber on pyboard...")
         try:
-            pyb.run_script(script, DevicePath("createstubs.py"))
+            pyb.run_script(create_stubs.getvalue(), DevicePath("createstubs.py"))
         except Exception as e:
             # TODO: Handle more usage cases
             self.log.error(f"Failed to execute script: {str(e)}", exception=e)
@@ -140,7 +149,7 @@ def _get_desc(name: str, cfg: dict):
             stub_path = next(out_dir.iterdir())
             self.log.info(f"Copied Stubs: $[{stub_path.name}]")
             stub_path = self.stubs.from_stubber(stub_path, out_dir)
-            stub = self.stubs.add(stub_path)
+            stub = self.stubs.add(str(stub_path))
         pyb.disconnect()
         self.log.success(f"Added {stub.name} to stubs!")
         return stub