create-gen: add --write flag to update __init__.py directly

Author: virtuald

robotpy_build/tool.py

@@ -9,6 +9,7 @@
 import subprocess
 import sys
 import types
+import typing
 import re
 from urllib.request import Request, urlopen
 from urllib.error import HTTPError
@@ -210,6 +211,9 @@ def add_subparser(cls, parent_parser, subparsers):
         )
         parser.add_argument("base", help="Ex: wpiutil")
         parser.add_argument("compiled", nargs="?", help="Ex: wpiutil._impl.wpiutil")
+        parser.add_argument(
+            "--write", "-w", action="store_true", help="Modify existing __init__.py"
+        )
         return parser
 
     def _rel(self, base: str, compiled: str) -> str:
@@ -220,16 +224,18 @@ def _rel(self, base: str, compiled: str) -> str:
         return f".{'.'.join(elems)}"
 
     def run(self, args):
+        self.create(args.base, args.compiled, args.write)
+
+    def create(self, base: str, compiled: typing.Optional[str], write: bool):
         # Runtime Dependency Check
         try:
             import black
         except:
             print("Error, The following module is required to run this tool: black")
             exit(1)
 
-        compiled = args.compiled
         if not compiled:
-            compiled = f"{args.base}._{args.base.split('.')[-1]}"
+            compiled = f"{base}._{base.split('.')[-1]}"
 
         # TODO: could probably generate this from parsed code, but seems hard
         ctx = {}
@@ -238,25 +244,57 @@ def run(self, args):
             if isinstance(ctx[k], types.ModuleType):
                 del ctx[k]
 
-        relimport = self._rel(args.base, compiled)
+        relimport = self._rel(base, compiled)
 
-        stmt_compiled = "" if not args.compiled else f" {args.compiled}"
+        stmt_compiled = "" if not compiled else f" {compiled}"
+        begin_stmt = f"# autogenerated by 'robotpy-build create-imports {base}"
 
         stmt = inspect.cleandoc(
             f"""
 
-            # autogenerated by 'robotpy-build create-imports {args.base}{stmt_compiled}'
+            {begin_stmt}{stmt_compiled}'
             from {relimport} import {','.join(sorted(ctx.keys()))}
             __all__ = ["{'", "'.join(sorted(ctx.keys()))}"]
         
         """
         )
 
-        print(
-            subprocess.check_output(
-                ["black", "-", "-q"], input=stmt.encode("utf-8")
-            ).decode("utf-8")
-        )
+        content = subprocess.check_output(
+            ["black", "-", "-q"], input=stmt.encode("utf-8")
+        ).decode("utf-8")
+
+        if write:
+            fctx = {}
+            exec(f"from {base} import __file__", {}, fctx)
+            fname = fctx["__file__"]
+
+            with open(fname) as fp:
+                fcontent = orig_content = fp.read()
+
+            # Find the beginning statement
+            idx = startidx = fcontent.find(begin_stmt)
+            if startidx != -1:
+                for to_find in ("from", "__all__", "[", "]", "\n"):
+                    idx = fcontent.find(to_find, idx)
+                    if idx == -1:
+                        startidx = -1
+                        break
+
+            if startidx == -1:
+                # If not present, just append and let the user figure it out
+                fcontent = fcontent + "\n" + content
+            else:
+                fcontent = fcontent[:startidx] + content + fcontent[idx + 1 :]
+
+            if fcontent != orig_content:
+                with open(fname, "w") as fp:
+                    fp.write(fcontent)
+                print("MOD", base)
+            else:
+                print("OK", base)
+
+        else:
+            print(content)
 
 
 class PlatformInfo: