Rewrite ErrorHandler

rcon/config.py

@@ -9,7 +9,7 @@ from os import getenv, name
 from pathlib import Path
 from typing import Iterable, NamedTuple, Optional, Union
 
-from rcon.exceptions import ConfigReadError
+from rcon.exceptions import ConfigReadError, UserAbort
 
 
 __all__ = ['CONFIG_FILES', 'LOG_FORMAT', 'SERVERS', 'Config', 'from_args']
@@ -93,7 +93,7 @@ def from_args(args: Namespace) -> Config:
             host, port, passwd = SERVERS[args.server]
         except KeyError:
             LOGGER.error('No such server: %s.', args.server)
-            raise ConfigReadError(2)
+            raise ConfigReadError() from None
 
     if passwd is None:
         try:
@@ -101,6 +101,6 @@ def from_args(args: Namespace) -> Config:
         except (KeyboardInterrupt, EOFError):
             print()
             LOGGER.error('Aborted by user.')
-            raise ConfigReadError(1)
+            raise UserAbort() from None
 
     return Config(host, port, passwd)

rcon/errorhandler.py

@@ -3,12 +3,25 @@
 from logging import Logger
 from socket import timeout
 
-from rcon.exceptions import ConfigReadError, RequestIdMismatch, WrongPassword
+from rcon.exceptions import ConfigReadError
+from rcon.exceptions import RequestIdMismatch
+from rcon.exceptions import UserAbort
+from rcon.exceptions import WrongPassword
 
 
 __all__ = ['ErrorHandler']
 
 
+ERRORS = {
+    UserAbort: 1,
+    ConfigReadError: 2,
+    ConnectionRefusedError: 3,
+    (TimeoutError, timeout): 4,
+    WrongPassword: 5,
+    RequestIdMismatch: 6
+}
+
+
 class ErrorHandler:
     """Handles common errors and exits."""
 
@@ -22,23 +35,11 @@ class ErrorHandler:
     def __enter__(self):
         return self
 
-    def __exit__(self, _, value, __):
+    def __exit__(self, _, value: Exception, __):
         """Checks for connection errors and exits respectively."""
-        if isinstance(value, ConnectionRefusedError):
-            self.logger.error('Connection refused.')
-            self.exit_code = 3
-        elif isinstance(value, (TimeoutError, timeout)):
-            self.logger.error('Connection timed out.')
-            self.exit_code = 4
-        elif isinstance(value, WrongPassword):
-            self.logger.error('Wrong password.')
-            self.exit_code = 5
-        elif isinstance(value, RequestIdMismatch):
-            self.logger.error('Session timed out.')
-            self.exit_code = 6
-        elif isinstance(value, ConfigReadError):
-            self.exit_code = value.exit_code
-        else:
-            return None
-
-        return True
+        for typ, exit_code in ERRORS.items():
+            if isinstance(value, typ):
+                self.exit_code = exit_code
+                return True
+
+        return None

rcon/exceptions.py

@@ -1,25 +1,30 @@
 """RCON exceptions."""
 
-__all__ = ['ConfigReadError', 'RequestIdMismatch', 'WrongPassword']
+__all__ = [
+    'ConfigReadError',
+    'RequestIdMismatch',
+    'UserAbort',
+    'WrongPassword'
+]
 
 
 class ConfigReadError(Exception):
     """Indicates an error while reading the configuration."""
 
-    def __init__(self, exit_code: int):
-        super().__init__()
-        self.exit_code = exit_code
-
 
 class RequestIdMismatch(Exception):
     """Indicates a mismatch in request IDs."""
 
     def __init__(self, sent: int, received: int):
-        """Sets the sent and received IDs."""
+        """Sets the IDs that have been sent and received."""
         super().__init__()
         self.sent = sent
         self.received = received
 
 
+class UserAbort(Exception):
+    """Indicates that a required action has been aborted by the user."""
+
+
 class WrongPassword(Exception):
     """Indicates a wrong password."""

rcon/rconclt.py

@@ -7,7 +7,6 @@ from pathlib import Path
 from rcon.client import Client
 from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args
 from rcon.errorhandler import ErrorHandler
-from rcon.exceptions import ConfigReadError
 
 
 __all__ = ['main']
@@ -40,24 +39,24 @@ def get_args() -> Namespace:
     return parser.parse_args()
 
 
-def main() -> int:
+def run() -> None:
     """Runs the RCON client."""
 
     args = get_args()
     basicConfig(format=LOG_FORMAT, level=DEBUG if args.debug else INFO)
+    host, port, passwd = from_args(args)
 
-    try:
-        host, port, passwd = from_args(args)
-    except ConfigReadError as cre:
-        return cre.exit_code
+    with Client(host, port, timeout=args.timeout) as client:
+        client.login(passwd)
+        text = client.run(args.command, *args.argument)
 
-    with ErrorHandler(LOGGER) as handler:
-        with Client(host, port, timeout=args.timeout) as client:
-            client.login(passwd)
-            text = client.run(args.command, *args.argument)
+    print(text, flush=True)
 
-    if handler.exit_code:
-        return handler.exit_code
 
-    print(text, flush=True)
-    return 0
+def main() -> int:
+    """Runs the main script with exceptions handled."""
+
+    with ErrorHandler(LOGGER) as handler:
+        run()
+
+    return handler.exit_code

rcon/rconshell.py

@@ -7,7 +7,6 @@ from pathlib import Path
 from rcon.config import CONFIG_FILES, LOG_FORMAT, from_args
 from rcon.console import PROMPT, rconcmd
 from rcon.errorhandler import ErrorHandler
-from rcon.exceptions import ConfigReadError
 from rcon.readline import CommandHistory
 
 
@@ -33,25 +32,25 @@ def get_args() -> Namespace:
     return parser.parse_args()
 
 
-def main() -> int:
+def run() -> None:
     """Runs the RCON shell."""
 
     args = get_args()
     basicConfig(level=INFO, format=LOG_FORMAT)
 
     if args.server:
-        try:
-            host, port, passwd = from_args(args)
-        except ConfigReadError as cre:
-            return cre.exit_code
+        host, port, passwd = from_args(args)
     else:
         host = port = passwd = None
 
-    with ErrorHandler(LOGGER) as handler:
-        with CommandHistory(LOGGER):
-            rconcmd(host, port, passwd, prompt=args.prompt)
+    with CommandHistory(LOGGER):
+        rconcmd(host, port, passwd, prompt=args.prompt)
+
 
-    if handler.exit_code:
-        return handler.exit_code
+def main() -> int:
+    """Runs the main script with exceptions handled."""
+
+    with ErrorHandler(LOGGER) as handler:
+        run()
 
-    return 0
+    return handler.exit_code