Try to please mypy

eX-Mech · Nov 10, 2022 · cd76b85 · cd76b85
1 parent 8882a23
commit cd76b85
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pymech/log.py b/pymech/log.py
@@ -15,6 +15,7 @@
 """
 import os
 import logging
+from typing import Union
 
 logger = logging.getLogger("pymech")
 
@@ -23,16 +24,15 @@
     # Set a nice colored output
     from rich.logging import RichHandler
 
-    handler = RichHandler()
-    logger.addHandler(handler)
+    handler: Union[RichHandler, logging.StreamHandler] = RichHandler()
 except ImportError:
     # No color available, use default config
     handler = logging.StreamHandler()
     formatter = logging.Formatter("%(levelname)s: %(message)s")
     handler.setFormatter(formatter)
-    logger.addHandler(handler)
     logger.info("Disabling coloured logs; if needed you should `pip install rich`.")
 
+logger.addHandler(handler)
 
 if bool(os.getenv("PYMECH_DEBUG")):
     logger.setLevel(logging.DEBUG)