From a823871039737c0a85674e5fd3334605d8f1825a Mon Sep 17 00:00:00 2001
From: Alek Ratzloff <alekratz@gmail.com>
Date: Fri, 3 Jun 2022 17:21:20 -0700
Subject: [PATCH] Add log level config option

* Log levels can now be set via the command line and the configuration
  file.
* ServerConfig.load() function takes a file-like object now, rather than
  a string

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>
---
 omnibot/__main__.py    | 30 +++++++++++++++++++++++++-----
 omnibot/config.py      | 18 ++++++++++++------
 tools/markov_import.py |  3 ++-
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/omnibot/__main__.py b/omnibot/__main__.py
index deeb90a..de25f8b 100644
--- a/omnibot/__main__.py
+++ b/omnibot/__main__.py
@@ -1,20 +1,40 @@
+import argparse
 import asyncio
 import logging
 
 from .config import ServerConfig
 from .bot import Bot
 
+parser = argparse.ArgumentParser(description="Run an IRC bot")
+parser.add_argument(
+    "-c", "--config", type=argparse.FileType("r"), default="config.toml"
+)
+parser.add_argument(
+    "--loglevel",
+    type=str,
+    choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+    required=False,
+)
+args = parser.parse_args()
+
+config = ServerConfig()
+config.load(args.config)
+
+# Set log level now
+loglevel = logging.INFO
+if args.loglevel:
+    loglevel = getattr(logging, args.loglevel)
+elif config.loglevel:
+    loglevel = getattr(logging, config.loglevel)
+
 logging.basicConfig(
-    level=logging.INFO,
+    level=loglevel,
     format="%(asctime)s - %(name)-12s - %(levelname)-8s - %(message)s",
 )
 log = logging.getLogger(__name__)
 
-log.debug("Loading config")
-config = ServerConfig()
-config.load("config.toml")
+log.info("Loaded config file %s", args.config)
 log.debug("Using configuration: %s", config)
-
 bot = Bot(config)
 
 try:
diff --git a/omnibot/config.py b/omnibot/config.py
index 9187166..38619b9 100644
--- a/omnibot/config.py
+++ b/omnibot/config.py
@@ -1,6 +1,6 @@
 import dataclasses
 from pathlib import Path
-from typing import Any, Mapping, Sequence, Set
+from typing import IO, Any, Mapping, Sequence, Set
 
 import toml
 
@@ -31,12 +31,10 @@ class ServerConfig:
     plugins: Sequence[PluginConfig] = dataclasses.field(default_factory=list)
     channels: Sequence[str] = dataclasses.field(default_factory=list)
     nick: str = "omnibot"
+    loglevel: str | None = None
 
-    def load(self, path: Path | str):
-        if isinstance(path, str):
-            path = Path(path)
-        with open(path) as fp:
-            obj = toml.load(fp)
+    def load(self, fp: IO[str]):
+        obj = toml.load(fp)
 
         if "server" not in obj:
             raise ConfigError("server", "must be present")
@@ -81,6 +79,14 @@ class ServerConfig:
                 raise ConfigError("nick", "must be a string")
             self.nick = obj["nick"]
 
+        if "loglevel" in obj:
+            if not isinstance(obj["loglevel"], str):
+                raise ConfigError("loglevel", "must be a string")
+            loglevels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
+            if obj["loglevel"] not in loglevels:
+                raise ConfigError("loglevel", "must be one of: " + " ".join(loglevels))
+            self.loglevel = obj["loglevel"]
+
     @property
     def all_channels(self) -> Set[str]:
         channels = set(self.channels)
diff --git a/tools/markov_import.py b/tools/markov_import.py
index 490f47a..e876876 100644
--- a/tools/markov_import.py
+++ b/tools/markov_import.py
@@ -23,7 +23,8 @@ async def main():
     channel = sys.argv[1]
     files = sys.argv[2:]
     server_config = ServerConfig()
-    server_config.load("config.toml")
+    with open("config.toml") as fp:
+        server_config.load(fp)
     # This only works on one plugin per config
     bot = Bot(server_config)
     plugin = [plugin for plugin in bot.plugins if isinstance(plugin, Markov)][0]