register extension with handler class

Author: manfred-kaiser

sifter/commands/fileinto.py

@@ -14,6 +14,7 @@
 class CommandFileInto(Command):
 
     HANDLER_ID = 'FILEINTO'
+    EXTENSION_NAME = 'fileinto'
     POSITIONAL_ARGS = [StringList(length=1)]
 
     def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions]:

sifter/commands/imap4flags.py

@@ -16,6 +16,7 @@
 class CommandSetFlag(Command):
 
     HANDLER_ID = 'SETFLAG'
+    EXTENSION_NAME = 'imap4flags'
     POSITIONAL_ARGS = [StringList()]
 
     def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions]:
@@ -29,6 +30,7 @@ def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions
 class CommandRemoveFlag(Command):
 
     HANDLER_ID = 'REMOVEFLAG'
+    EXTENSION_NAME = 'imap4flags'
     POSITIONAL_ARGS = [StringList()]
 
     def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions]:
@@ -42,6 +44,7 @@ def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions
 class CommandAddFlag(Command):
 
     HANDLER_ID = 'ADDFLAG'
+    EXTENSION_NAME = 'imap4flags'
     POSITIONAL_ARGS = [StringList()]
 
     def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions]:

sifter/commands/notify.py

@@ -16,6 +16,7 @@
 class CommandNotify(Command):
 
     HANDLER_ID = 'NOTIFY'
+    EXTENSION_NAME = 'enotify'
     TAGGED_ARGS = {
         'from': Tag('FROM', (StringList(1),)),
         'importance': Tag('IMPORTANCE', (StringList(1),)),

sifter/commands/reject.py

@@ -14,6 +14,7 @@
 class CommandReject(Command):
 
     HANDLER_ID: Text = 'REJECT'
+    EXTENSION_NAME = 'reject'
     POSITIONAL_ARGS = [StringList()]
 
     def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions]:
@@ -26,3 +27,4 @@ def evaluate(self, message: Message, state: EvaluationState) -> Optional[Actions
 class CommandEReject(CommandReject):
 
     HANDLER_ID: Text = 'EREJECT'
+    EXTENSION_NAME = 'ereject'

sifter/commands/variables.py

@@ -19,6 +19,7 @@
 class CommandSet(Command):
 
     HANDLER_ID = 'SET'
+    EXTENSION_NAME = 'variables'
     TAGGED_ARGS = {
         'lower': Tag('LOWER'),
         'upper': Tag('UPPER'),

sifter/extensions/__init__.py

@@ -1,5 +1,6 @@
 from typing import (
     TYPE_CHECKING,
+    cast,
     Dict,
     List,
     Text,
@@ -17,6 +18,7 @@
 from sifter.grammar.test import Test
 if TYPE_CHECKING:
     from sifter.grammar.tag import Tag
+    from sifter.grammar.sieveobject import SieveObject
 
 
 class ExtensionRegistry():
@@ -27,38 +29,32 @@ class ExtensionRegistry():
             Text,
             Union[
                 bool,
-                Type['Comparator'],
-                Type['Rule'],
-                Type['NotificationMethod']
+                Type['SieveObject']
             ]
         ]
     ] = {}
     DEFAULT_EXTENSION: List[Text] = [
-        'body',
         'comparator-i;ascii-casemap',
-        'comparator-i;octet',
-        'enotify',
-        'fileinto',
-        'imap4flags',
-        'regex',
-        'reject',
-        'ereject',
-        'variables',
+        'comparator-i;octet'
+        'regex'
     ]
 
     def __init__(self) -> None:
         for extension_name in self.DEFAULT_EXTENSION:
             self.register_extension(extension_name)
 
         for entry_point in pkg_resources.iter_entry_points('sifter_extensions'):
-            self.register_handler(entry_point.load())
+            sifter_extension_cls = cast(Type['SieveObject'], entry_point.load())
+            self.register_handler(sifter_extension_cls)
+            if sifter_extension_cls.EXTENSION_NAME is not None:
+                self.register_extension(sifter_extension_cls.EXTENSION_NAME)
 
     @classmethod
     def register_extension(cls, extension_name: Text) -> None:
         cls.register('extension', extension_name, True)
 
     @classmethod
-    def register_handler(cls, ext_cls: Union[Type['Comparator'], Type['Rule']]) -> None:
+    def register_handler(cls, ext_cls: Type['SieveObject']) -> None:
         cls.register(ext_cls.handler_type(), ext_cls.handler_id(), ext_cls)
 
     @classmethod
@@ -100,18 +96,18 @@ def register(
         cls,
         handler_type: Text,
         handler_id: Text,
-        value: Union[bool, Type['Comparator'], Type['Rule']]
+        value: Union[bool, Type['SieveObject']]
     ) -> None:
         cls._HANDLERS_MAP.setdefault(handler_type, {})[handler_id] = value
 
     @classmethod
     def unregister(
         cls, handler_type: Text, handler_id: Text
-    ) -> Optional[Union[bool, Type[NotificationMethod], Type['Comparator'], Type['Rule']]]:
+    ) -> Optional[Union[bool, Type['SieveObject']]]:
         return cls._HANDLERS_MAP.get(handler_type, {}).pop(handler_id, None)
 
     @classmethod
     def get(
         cls, handler_type: Text, handler_id: Text
-    ) -> Optional[Union[bool, Type[NotificationMethod], Type['Comparator'], Type['Rule']]]:
+    ) -> Optional[Union[bool, Type['SieveObject']]]:
         return cls._HANDLERS_MAP.get(handler_type, {}).get(handler_id, None)

sifter/grammar/sieveobject.py

@@ -7,6 +7,7 @@ class SieveObject():
 
     HANDLER_TYPE: Optional[Text] = None
     HANDLER_ID: Optional[Text] = None
+    EXTENSION_NAME: Optional[Text] = None
 
     @classmethod
     def handler_type(cls) -> Text:

sifter/tests/body.py

@@ -24,6 +24,7 @@
 class TestBody(Test):
 
     HANDLER_ID = 'BODY'
+    EXTENSION_NAME = 'body'
     TAGGED_ARGS = {
         'comparator': Comparator(),
         'match_type': MatchType(),