refactor stuff

kos_sim/__init__.py

@@ -2,10 +2,6 @@
 
 import logging
 
-import colorlogging
-
-colorlogging.configure()
-
 logger = logging.getLogger("kos_sim")
 
 __version__ = "0.0.1"

kos_sim/server.py

@@ -5,6 +5,7 @@
 import time
 from concurrent import futures
 
+import colorlogging
 import grpc
 from kos_protos import actuator_pb2_grpc, imu_pb2_grpc, sim_pb2_grpc
 
@@ -93,11 +94,17 @@ def serve(model_path: str, config_path: str | None = None, port: int = 50051) ->
     server.start()
 
 
-if __name__ == "__main__":
+def run_server() -> None:
     parser = argparse.ArgumentParser(description="Start the simulation gRPC server.")
     parser.add_argument("--model-path", type=str, required=True, help="Path to MuJoCo XML model file")
     parser.add_argument("--port", type=int, default=50051, help="Port to listen on")
     parser.add_argument("--config-path", type=str, default=None, help="Path to config file")
 
+    colorlogging.configure()
+
     args = parser.parse_args()
     serve(args.model_path, args.config_path, args.port)
+
+
+if __name__ == "__main__":
+    run_server()

kos_sim/simulator.py

@@ -5,6 +5,7 @@
 import threading
 from pathlib import Path
 
+import colorlogging
 import mujoco
 import mujoco_viewer
 import numpy as np
@@ -237,6 +238,8 @@ async def main() -> None:
     parser.add_argument("--speed", type=float, default=1.0, help="Simulation speed multiplier")
     parser.add_argument("--no-render", action="store_true", help="Disable rendering")
 
+    colorlogging.configure()
+
     args = parser.parse_args()
     await test_simulation_adhoc(
         args.model_name,

kos_sim/test_client.py

@@ -94,7 +94,7 @@ def test_sim_service(host: str = "localhost", port: int = 50051) -> None:
         logger.error("Test failed: %s", e)
 
 
-if __name__ == "__main__":
+def main() -> None:
     parser = argparse.ArgumentParser(description="Test the simulation server with actuator commands.")
     parser.add_argument("--host", type=str, default="localhost", help="Server host")
     parser.add_argument("--port", type=int, default=50051, help="Server port")
@@ -105,3 +105,7 @@ def test_sim_service(host: str = "localhost", port: int = 50051) -> None:
         test_actuator_commands(args.host, args.port)
     else:
         test_sim_service(args.host, args.port)
+
+
+if __name__ == "__main__":
+    main()