Tools for handling env var overrides

skel/Desktop/simulated-drone/sim_drone.py

@@ -1,57 +1,104 @@
 #!/bin/env python3
-import typer, os, sys
+import typer, os, sys, contextlib
 from loguru import logger
 import sh
 import atexit
 
 logger.remove()
-logger.add(sys.stdout, format="<green>{time}</green> <level>{level}</level> {extra} {message}")
+logger.add(
+    sys.stdout, format="<green>{time}</green> <level>{level}</level> {extra} {message}"
+)
 
-px4Path = os.environ.get("SPIRI_SIM_PX4_PATH","/opt/spiri-sdk/PX4-Autopilot/build/px4_sitl_default/bin/px4")
+px4Path = os.environ.get(
+    "SPIRI_SIM_PX4_PATH", "/opt/spiri-sdk/PX4-Autopilot/build/px4_sitl_default/bin/px4"
+)
 logger.info(f"SPIRI_SIM_PX4_PATH={px4Path}")
 px4 = sh.Command(px4Path)
 
 app = typer.Typer()
 
-class outputLogger():
+
+class outputLogger:
     """
     Logs command output to loguru
     """
+
     def __init__(self, name):
-        self.name=name
+        self.name = name
+
     def __call__(self, message):
         with logger.contextualize(cmd=self.name):
-            if message.endswith('\n'):
-              message = message[:-1]
-            #ToDo, this doesn't work because the output is coloured
-            if message.startswith('INFO'):
-                message=message.lstrip("INFO")
+            if message.endswith("\n"):
+                message = message[:-1]
+            # ToDo, this doesn't work because the output is coloured
+            if message.startswith("INFO"):
+                message = message.lstrip("INFO")
                 logger.info(message)
-            elif message.startswith('WARN'): 
-                message=message.lstrip("WARN")
+            elif message.startswith("WARN"):
+                message = message.lstrip("WARN")
                 logger.warning(message)
+            elif message.startswith("ERROR"):
+                message = message.lstrip("ERROR")
+                logger.error(message)
             else:
                 logger.info(message)
 
+
+
+@contextlib.contextmanager
+def modified_environ(*remove, **update):
+    """
+    Temporarily updates the ``os.environ`` dictionary in-place.
+
+    The ``os.environ`` dictionary is updated in-place so that the modification
+    is sure to work in all situations.
+
+    :param remove: Environment variables to remove.
+    :param update: Dictionary of environment variables and values to add/update.
+    """
+    env = os.environ
+    update = update or {}
+    remove = remove or []
+
+    # List of environment variables being updated or removed.
+    stomped = (set(update.keys()) | set(remove)) & set(env.keys())
+    # Environment variables and values to restore on exit.
+    update_after = {k: env[k] for k in stomped}
+    # Environment variables and values to remove on exit.
+    remove_after = frozenset(k for k in update if k not in env)
+
+    try:
+        env.update(update)
+        [env.pop(k, None) for k in remove]
+        yield
+    finally:
+        env.update(update_after)
+        [env.pop(k) for k in remove_after]
+
+
 @app.command()
-def start(sys_id: int =1):
+def start(sys_id: int = 1):
     """Starts the simulated drone with a given sys_id,
     each drone must have it's own unique ID.
     """
-    instance_id = sys_id-1 #PX4 will add 1 to the instance ID to get sysID
-    px4Instance = px4(i=instance_id, 
-            _out=outputLogger("px4"), _bg=True)
-    atexit.register(lambda: px4Instance.kill())
-    docker_stack = sh.docker.compose(
-            "--project-name", f"simulated_robot_{sys_id}",
-            "up",
-            _out=outputLogger("docker_stack"),
-            _bg=True
-            )
-    atexit.register(lambda: docker_stack.kill())
-    px4Instance.wait()
-    docker_stack.wait()
-
+    with modified_environ(
+            PX4_SIM_MODEL=os.environ.get("PX4_SIM_MODEL") or "gz_x500",
+            PX4_GZ_MODEL_POSE=os.environ.get("PX4_GZ_MODEL_POSE") or f"0,{sys_id-1}" 
+        ):
+        instance_id = sys_id - 1  # PX4 will add 1 to the instance ID to get sysID
+        px4Instance = px4(i=instance_id, _out=outputLogger("px4"), _bg=True)
+        atexit.register(lambda: px4Instance.kill())
+        logger.info("Starting drone stack, this may take some time")
+        # docker_stack = sh.docker.compose(
+        #         "--project-name", f"simulated_robot_{sys_id}",
+        #         "up",
+        #         _out=outputLogger("docker_stack"),
+        #         _err=sys.stderr,
+        #         _bg=True
+        #         )
+        # atexit.register(lambda: docker_stack.kill())
+        px4Instance.wait()
+        # docker_stack.wait()
 
 
 if __name__ == "__main__":