More typing

plugins/python-dspam-plugin-osb/src/dspam_plugin_osb/osb.py

@@ -23,10 +23,10 @@ class OsbTokenizer(Tokenizer):
     SPARSE_WINDOW_SIZE = 5
     """The size of the sliding window for generating OSB tokens."""
 
-    TOKEN_SEPARATOR: str = "+"  # noqa: S105
+    TOKEN_SEPARATOR = "+"  # noqa: S105
     """The separator to use between terms in OSB tokens."""
 
-    OSB_TOKEN_PLACEHOLDER = "#"  # noqa: S105
+    TOKEN_PLACEHOLDER = "#"  # noqa: S105
     """The placeholder for skipped terms in an OSB token"""
 
     def __init__(self, settings: TokenizerSettings) -> None:
@@ -94,7 +94,7 @@ def get_osb_tokens(self, word_tokens: TokenList) -> Iterator[Token]:
                 elif osb_token_terms:
                     # Append placeholders only when we already have other terms
                     logger.debug(f"Appending placeholder for term {term}")
-                    osb_token_terms.append(self.OSB_TOKEN_PLACEHOLDER)
+                    osb_token_terms.append(self.TOKEN_PLACEHOLDER)
 
             osb_token = self.TOKEN_SEPARATOR.join(osb_token_terms)
             logger.debug(f"Generated OSB token: {osb_token}")

src/dspam/__init__.py

@@ -1,9 +1 @@
 # SPDX-License-Identifier: BSD-3-Clause
-
-
-"""Some constants used throughout the application"""
-
-# TODO: change these into an Enum class or whatnot
-IS_INNOCENT: str = "innocent"
-IS_SPAM: str = "spam"
-IS_UNKNOWN: str = "unknown"

src/dspam/classify.py

@@ -4,9 +4,9 @@
 import random
 from abc import ABC, abstractmethod
 
-from dspam import IS_INNOCENT, IS_SPAM, IS_UNKNOWN
 from dspam.settings import ClassifierSettings
 from dspam.storage import Storage
+from dspam.types import Classification, TokenList
 
 logger = logging.getLogger(__name__)
 
@@ -19,13 +19,13 @@ def __init__(self, settings: ClassifierSettings, storage: Storage) -> None:
         self.storage = storage
 
     @abstractmethod
-    async def __call__(self, tokens: list[str]) -> str:
+    async def __call__(self, tokens: list[str]) -> Classification:
         pass  # pragma: no cover
 
     def __str__(self) -> str:
         return f"{self.__class__.__name__}(API_VERSION={self.API_VERSION})"  # pragma: no cover
 
-    async def log_debug_tokens(self, tokens: list[str], verdict: str, debug_token_count: int = 0) -> None:
+    async def log_debug_tokens(self, tokens: TokenList, verdict: Classification, debug_token_count: int = 0) -> None:
         """
         Utility method to log some tokens for debugging purposes.
 
@@ -52,11 +52,11 @@ class DummyClassifier(Classifier):
 
     API_VERSION: str = "1.0"
 
-    async def __call__(self, tokens: list[str]) -> str:
+    async def __call__(self, tokens: TokenList) -> Classification:
         for token in tokens:
             if "spam" in token:
-                return IS_SPAM
-        return IS_INNOCENT
+                return Classification.SPAM
+        return Classification.INNOCENT
 
 
 class SimpleClassifier(Classifier):
@@ -73,7 +73,7 @@ class SimpleClassifier(Classifier):
     The tokens are logged at level DEBUG. For each verdict (innocent, ham), this number of tokens will be logged.
     """
 
-    async def __call__(self, tokens: list[str]) -> str:
+    async def __call__(self, tokens: TokenList) -> Classification:
         spam_tokens = []
         innocent_tokens = []
         unknown_tokens = []
@@ -90,16 +90,16 @@ async def __call__(self, tokens: list[str]) -> str:
                 innocent_tokens.append(token)
 
         # Log some tokens for debugging purposes
-        await self.log_debug_tokens(spam_tokens, IS_SPAM, self.DEBUG_TOKEN_COUNT)
-        await self.log_debug_tokens(innocent_tokens, IS_INNOCENT, self.DEBUG_TOKEN_COUNT)
-        await self.log_debug_tokens(unknown_tokens, IS_UNKNOWN, self.DEBUG_TOKEN_COUNT)
+        await self.log_debug_tokens(spam_tokens, Classification.SPAM, self.DEBUG_TOKEN_COUNT)
+        await self.log_debug_tokens(innocent_tokens, Classification.INNOCENT, self.DEBUG_TOKEN_COUNT)
+        await self.log_debug_tokens(unknown_tokens, Classification.UNKNOWN, self.DEBUG_TOKEN_COUNT)
 
         spam_count = len(spam_tokens)
         innocent_count = len(innocent_tokens)
         unknown_count = len(unknown_tokens)
 
         logger.info(f"Classifier token results: spam={spam_count}, innocent={innocent_count}, unknown={unknown_count}")
         if spam_count > innocent_count:
-            return IS_SPAM
+            return Classification.SPAM
         else:
-            return IS_INNOCENT
+            return Classification.INNOCENT

src/dspam/cli.py

@@ -19,7 +19,7 @@
 from dspam.main import classify, train
 from dspam.plugins import PluginManager
 from dspam.settings import Settings
-from dspam.types import Verdict
+from dspam.types import Classification
 
 cli = App(help=__doc__, version=version("python-dspam"))
 
@@ -49,7 +49,7 @@ async def train_from_file(
         ),
     ],
     classification: Annotated[
-        Verdict,
+        Classification,
         Parameter(alias="-c", help="Classification of the file (innocent or spam)"),
     ],
 ) -> None:

src/dspam/di.py

@@ -17,6 +17,7 @@
 from dspam.storage import Storage, get_storage_root
 from dspam.tokenize import Tokenizer
 from dspam.train import Trainer
+from dspam.types import PluginGroup
 
 logger = logging.getLogger(__name__)
 
@@ -32,7 +33,7 @@ def settings_factory(context: ActivationScope) -> Settings:
     tmp_settings = Settings()
 
     settings_kwargs = {}
-    for group_name in pm.GROUPS:
+    for group_name in PluginGroup:
         group_settings = getattr(tmp_settings, group_name)
         plugin_name = group_settings.plugin
         plugin_settings = pm.get_plugin_settings(group_name, plugin_name)
@@ -47,7 +48,7 @@ def parser_factory(context: ActivationScope) -> Parser:
     pm = context.provider.get(PluginManager)
     settings = context.provider.get(Settings)
 
-    parser_class: type[Parser] = pm.get_plugin(pm.PARSER, settings.parser.plugin)
+    parser_class: type[Parser] = pm.get_plugin(PluginGroup.PARSER, settings.parser.plugin)
     parser_instance = parser_class(settings=settings.parser)
     logger.debug(f"Initialized parser: {parser_instance}")
     return parser_instance
@@ -57,7 +58,7 @@ def tokenizer_factory(context: ActivationScope) -> Tokenizer:
     """Factory for DI-supplied Tokenizer instances."""
     pm = context.provider.get(PluginManager)
     settings = context.provider.get(Settings)
-    tokenizer_class: type[Tokenizer] = pm.get_plugin(pm.TOKENIZER, settings.tokenizer.plugin)
+    tokenizer_class: type[Tokenizer] = pm.get_plugin(PluginGroup.TOKENIZER, settings.tokenizer.plugin)
 
     tokenizer_instance = tokenizer_class(settings=settings.tokenizer)
     logger.debug(f"Initialized tokenizer: {tokenizer_instance}")
@@ -68,7 +69,7 @@ def storage_factory(context: ActivationScope) -> Storage:
     """Factory for DI-supplied Storage instances."""
     pm = context.provider.get(PluginManager)
     settings = context.provider.get(Settings)
-    storage_class: type[Storage] = pm.get_plugin(pm.STORAGE, settings.storage.plugin)
+    storage_class: type[Storage] = pm.get_plugin(PluginGroup.STORAGE, settings.storage.plugin)
 
     storage_root = get_storage_root()
     storage_instance = storage_class(settings=settings.storage, storage_root=storage_root)
@@ -80,7 +81,7 @@ def classify_factory(context: ActivationScope) -> Classifier:
     """Factory for DI-supplied Classifier instances."""
     pm = context.provider.get(PluginManager)
     settings = context.provider.get(Settings)
-    classification_class: type[Classifier] = pm.get_plugin(pm.CLASSIFIER, settings.classifier.plugin)
+    classification_class: type[Classifier] = pm.get_plugin(PluginGroup.CLASSIFIER, settings.classifier.plugin)
     storage = context.provider.get(Storage)
 
     classifier_instance = classification_class(settings.classifier, storage)
@@ -92,7 +93,7 @@ def trainer_factory(context: ActivationScope) -> Trainer:
     """Factory for DI-supplied Trainer instances."""
     pm = context.provider.get(PluginManager)
     settings = context.provider.get(Settings)
-    trainer_class: type[Trainer] = pm.get_plugin(pm.TRAINER, settings.trainer.plugin)
+    trainer_class: type[Trainer] = pm.get_plugin(PluginGroup.TRAINER, settings.trainer.plugin)
     storage = context.provider.get(Storage)
 
     trainer_instance = trainer_class(settings.trainer, storage)

src/dspam/main.py

@@ -9,7 +9,7 @@
 from dspam.parse import Parser
 from dspam.tokenize import Tokenizer
 from dspam.train import Trainer
-from dspam.types import Verdict
+from dspam.types import Classification
 
 logger = logging.getLogger(__name__)
 
@@ -35,7 +35,7 @@ async def classify(file_path: Path) -> None:
     print(f"Classification result for {file_path}: {classification}")
 
 
-async def train(file_path: Path, classification: Verdict) -> None:
+async def train(file_path: Path, classification: Classification) -> None:
     """Train the classifier with the given file and classification."""
     tokens = await _parse_and_tokenize(file_path)
 

src/dspam/plugins.py

@@ -20,12 +20,13 @@
 import importlib.metadata
 import inspect
 import logging
-from collections.abc import Generator
+from collections.abc import Iterator
 from dataclasses import dataclass
 
 from pydantic_settings import BaseSettings
 
 from dspam.exceptions import DspamPluginNotFound
+from dspam.types import PluginGroup
 
 logger = logging.getLogger(__name__)
 
@@ -58,26 +59,11 @@ class PluginInfo:
 
 
 class PluginManager:
-    # Plugin groups
-    PARSER: str = "parser"
-    TOKENIZER: str = "tokenizer"
-    CLASSIFIER: str = "classifier"
-    TRAINER: str = "trainer"
-    STORAGE: str = "storage"
-
-    GROUPS: list[str] = [
-        PARSER,
-        TOKENIZER,
-        CLASSIFIER,
-        TRAINER,
-        STORAGE,
-    ]
-
-    PLUGIN_ENTRY_POINTS = {f"dspam.{group}": group for group in GROUPS}
+    PLUGIN_ENTRY_POINTS = {f"dspam.{group}": group for group in PluginGroup}
     """Supported entry point groups for registering plugins: dspam:<group name>."""
 
     def __init__(self) -> None:
-        self.plugins: dict[str, dict[str, type]] = {group: {} for group in self.GROUPS}
+        self.plugins: dict[str, dict[str, type]] = {group: {} for group in PluginGroup}
 
     def load_all_plugins(self) -> None:
         """
@@ -92,12 +78,12 @@ def load_builtin_plugins(self) -> None:
         classify_module = importlib.import_module("dspam.classify")
         training_module = importlib.import_module("dspam.train")
         storage_module = importlib.import_module("dspam.storage")
-        self.plugins[self.PARSER]["plaintext"] = parse_module.PlainTextParser
-        self.plugins[self.TOKENIZER]["word"] = tokenize_module.WordTokenizer
-        self.plugins[self.CLASSIFIER]["simple"] = classify_module.SimpleClassifier
-        self.plugins[self.CLASSIFIER]["dummy"] = classify_module.DummyClassifier
-        self.plugins[self.TRAINER]["simple"] = training_module.SimpleTrainer
-        self.plugins[self.STORAGE]["json"] = storage_module.JSONStorage
+        self.plugins[PluginGroup.PARSER]["plaintext"] = parse_module.PlainTextParser
+        self.plugins[PluginGroup.TOKENIZER]["word"] = tokenize_module.WordTokenizer
+        self.plugins[PluginGroup.CLASSIFIER]["simple"] = classify_module.SimpleClassifier
+        self.plugins[PluginGroup.CLASSIFIER]["dummy"] = classify_module.DummyClassifier
+        self.plugins[PluginGroup.TRAINER]["simple"] = training_module.SimpleTrainer
+        self.plugins[PluginGroup.STORAGE]["json"] = storage_module.JSONStorage
 
         plugin_names = [f"{p.group}:{p.name}" for p in self.list_plugins()]
         logger.debug(f"Loaded built-in plugins: {', '.join(plugin_names)}")
@@ -112,7 +98,7 @@ def load_entrypoint_plugins(self) -> None:
                 except Exception as err:
                     logger.exception(f"Error loading plugin {group}:{entry_point.name}", exc_info=err)
 
-    def list_plugins(self) -> Generator[PluginInfo]:
+    def list_plugins(self) -> Iterator[PluginInfo]:
         """
         List all loaded plugins.
 
@@ -133,7 +119,7 @@ def list_plugins(self) -> Generator[PluginInfo]:
                     api_version=api_version,
                 )
 
-    def get_plugin(self, group_name: str, plugin_name: str) -> type:
+    def get_plugin(self, group_name: PluginGroup, plugin_name: str) -> type:
         """
         Get the plugin class for the specified group and plugin name.
 
@@ -147,7 +133,7 @@ def get_plugin(self, group_name: str, plugin_name: str) -> type:
             raise DspamPluginNotFound(f"Plugin {group_name}.{plugin_name} not found")
         return plugin
 
-    def get_plugin_settings(self, group_name: str, plugin_name: str) -> type[BaseSettings] | None:
+    def get_plugin_settings(self, group_name: PluginGroup, plugin_name: str) -> type[BaseSettings] | None:
         """
         Get the settings class from the plugin module, if available, and return it.
 

src/dspam/settings.py

@@ -2,9 +2,8 @@
 
 import os
 import pathlib
-from functools import cached_property
 
-from pydantic import BaseModel, computed_field
+from pydantic import BaseModel
 from pydantic_settings import (
     BaseSettings,
     PydanticBaseSettingsSource,
@@ -13,7 +12,7 @@
 )
 
 from dspam.plugins import PluginManager
-from dspam.types import LogLevel
+from dspam.types import LogLevel, PluginGroup
 
 
 def get_config_root() -> pathlib.Path:
@@ -30,7 +29,7 @@ def get_config_file() -> pathlib.Path | None:
     return None
 
 
-def get_plugin_settings(group_name: str, plugin_name: str) -> type[BaseSettings] | None:
+def get_plugin_settings(group_name: PluginGroup, plugin_name: str) -> type[BaseSettings] | None:
     from dspam.di import provider
 
     pm = provider.get(PluginManager)
@@ -80,45 +79,21 @@ class TokenizerSettings(BaseModel):
     delimiters: str = " .,!?;:\"''@()[]{}<>=*/\\"
     """The list of delimiters that a tokenizer uses to separate content into tokens."""
 
-    @computed_field(exclude_if=lambda x: x is None)  # type: ignore[prop-decorator]
-    @cached_property
-    def plugin_settings(self) -> type[BaseSettings] | None:
-        settings = get_plugin_settings(PluginManager.TOKENIZER, self.plugin)
-        return settings
-
 
 class StorageSettings(BaseModel):
     plugin: str = "json"
     """The plugin to use for storing data."""
 
-    @computed_field(exclude_if=lambda x: x is None)  # type: ignore[prop-decorator]
-    @cached_property
-    def plugin_settings(self) -> type[BaseSettings] | None:
-        settings = get_plugin_settings(PluginManager.STORAGE, self.plugin)
-        return settings
-
 
 class ClassifierSettings(BaseModel):
     plugin: str = "simple"
     """The plugin to use for classification."""
 
-    @computed_field(exclude_if=lambda x: x is None)  # type: ignore[prop-decorator]
-    @cached_property
-    def plugin_settings(self) -> type[BaseSettings] | None:
-        settings = get_plugin_settings(PluginManager.CLASSIFIER, self.plugin)
-        return settings
-
 
 class TrainerSettings(BaseModel):
     plugin: str = "simple"
     """The plugin to use for training."""
 
-    @computed_field(exclude_if=lambda x: x is None)  # type: ignore[prop-decorator]
-    @cached_property
-    def plugin_settings(self) -> type[BaseSettings] | None:
-        settings = get_plugin_settings(PluginManager.TRAINER, self.plugin)
-        return settings
-
 
 class Settings(BaseParserSettings, BaseSettings):
     # TODO: Remove type ignore after release of: https://github.com/pydantic/pydantic-settings/pull/882