diff --git a/skel/Desktop/simulated-drone/sim_drone.py b/skel/Desktop/simulated-drone/sim_drone.py index 4b2ee50..9d8b235 100755 --- a/skel/Desktop/simulated-drone/sim_drone.py +++ b/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="{time} {level} {extra} {message}") +logger.add( + sys.stdout, format="{time} {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__":