Repository: aio-libs/aiosmtpd
Pull request #6 (2016-08-12 11:59:10)

AsyncMessage handler

kozzztik: I dont really sure, that you must force loop as handler kwarg. User may create on object with loop parameter and provide object method where loop can be accessible through self variable.
sirkonst: In this case, problems can arise when instance of AsyncMessage class created in one thread and smtpd server runn in other. And this threads has different loops.  It's real case if use aiosmtpd.controller.Controller like http://aiosmtpd.readthedocs.io/en/latest/aiosmtpd/docs/controller.html
warsaw: I think you're right about that.  Given the current structure of the code, the best way to handle that would be to pass the `AsyncMessage` *class* in as the handler, and override the `factory()` function in a subclass of the `Controller` to instantiate it.  E.g. (untested):

```python

class AsyncMessageController(Controller):
    def factory(self):
        return SMTP(self.handler(self.loop))

my_controller = AsyncMessageController(AsyncMessage)
```
warsaw: Thanks for the contribution!  Other than the review comments, I would like to ask if you could add some documentation for the `AsyncMessage` handler.  I wish we had better documentation for the base `Message` class too.  Would you consider adding a .rst file for documenting the `AsyncMessage` class and why you'd choose it over say the base `Message` class?

Patch:
diff --git a/aiosmtpd/handlers.py b/aiosmtpd/handlers.py
index 04dfee0..d48add3 100644
--- a/aiosmtpd/handlers.py
+++ b/aiosmtpd/handlers.py
@@ -6,8 +6,8 @@ Pass in an instance of one of these classes, or derive your own, to provide
 your own handling of messages.  Implement only the methods you care about.
 """
 
-
 import sys
+import asyncio
 import logging
 import mailbox
 import smtplib
@@ -136,7 +136,7 @@ class Message:
     def __init__(self, message_class=None):
         self.message_class = message_class
 
-    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+    def prepare_message(self, peer, mailfrom, rcpttos, data, **kws):
         # If the server was created with decode_data True, then data will be a
         # str, otherwise it will be bytes.
         if isinstance(data, bytes):
@@ -148,12 +148,33 @@ class Message:
         message['X-Peer'] = str(peer)
         message['X-MailFrom'] = mailfrom
         message['X-RcptTos'] = COMMASPACE.join(rcpttos)
+
+        return message
+
+    def process_message(self, peer, mailfrom, rcpttos, data, **kws):
+        message = self.prepare_message(peer, mailfrom, rcpttos, data, **kws)
         self.handle_message(message)
 
     def handle_message(self, message):
         raise NotImplementedError                   # pragma: no cover
 
 
+@public
+class AsyncMessage(Message):
+    def __init__(self, message_class=None, *, loop=None):
+        super().__init__(message_class)
+        self.loop = loop or asyncio.get_event_loop()
+
+    @asyncio.coroutine
+    def process_message(self, peer, mailfrom, rcpttos, data, *, loop, **kws):
+        message = self.prepare_message(peer, mailfrom, rcpttos, data, **kws)
+        yield from self.handle_message(message, loop=loop)
+
+    @asyncio.coroutine
+    def handle_message(self, message, *, loop):
+        raise NotImplementedError                   # pragma: no cover
+
+
 @public
 class Mailbox(Message):
     def __init__(self, mail_dir, message_class=None):
diff --git a/aiosmtpd/smtp.py b/aiosmtpd/smtp.py
index bcaccfc..2d1e134 100644
--- a/aiosmtpd/smtp.py
+++ b/aiosmtpd/smtp.py
@@ -438,7 +438,12 @@ class SMTP(asyncio.StreamReaderProtocol):
                 'mail_options': self.mail_options,
                 'rcpt_options': self.rcpt_options,
             }
-        status = self.event_handler.process_message(*args, **kwargs)
+        kwargs.update({'loop': self.loop})
+        if asyncio.iscoroutinefunction(self.event_handler.process_message):
+            status = yield from self.event_handler.process_message(
+                *args, **kwargs)
+        else:
+            status = self.event_handler.process_message(*args, **kwargs)
         self._set_post_data_state()
         if status:
             yield from self.push(status)


Test patch:
diff --git a/aiosmtpd/tests/test_handlers.py b/aiosmtpd/tests/test_handlers.py
index 989b0c0..9253565 100644
--- a/aiosmtpd/tests/test_handlers.py
+++ b/aiosmtpd/tests/test_handlers.py
@@ -1,9 +1,11 @@
 import os
 import sys
+import asyncio
 import unittest
 
 from aiosmtpd.controller import Controller
-from aiosmtpd.handlers import Debugging, Mailbox, Message, Proxy, Sink
+from aiosmtpd.handlers import (
+    AsyncMessage, Debugging, Mailbox, Message, Proxy, Sink)
 from aiosmtpd.smtp import SMTP as Server
 from contextlib import ExitStack
 from io import StringIO
@@ -186,6 +188,64 @@ Testing
         self.assertEqual(self.handled_message['X-RcptTos'], 'bart@example.com')
 
 
+class TestAsyncMessage(unittest.TestCase):
+    def setUp(self):
+        self.handled_message = None
+
+        class MessageHandler(AsyncMessage):
+            @asyncio.coroutine
+            def handle_message(handler_self, message, loop):
+                self.handled_message = message
+
+        self.handler = MessageHandler()
+
+    def test_message(self):
+        # In this test, the message data comes in as bytes.
+        controller = Controller(self.handler)
+        controller.start()
+        self.addCleanup(controller.stop)
+        with SMTP(controller.hostname, controller.port) as client:
+            client.sendmail('anne@example.com', ['bart@example.com'], """\
+From: Anne Person <anne@example.com>
+To: Bart Person <bart@example.com>
+Subject: A test
+Message-ID: <ant>
+
+Testing
+""")
+        self.assertEqual(self.handled_message['subject'], 'A test')
+        self.assertEqual(self.handled_message['message-id'], '<ant>')
+        self.assertIsNotNone(self.handled_message['X-Peer'])
+        self.assertEqual(
+            self.handled_message['X-MailFrom'], 'anne@example.com')
+        self.assertEqual(self.handled_message['X-RcptTos'], 'bart@example.com')
+
+    def test_message_decoded(self):
+        # With a server that decodes the data, the messages come in as
+        # strings.  There's no difference in the message seen by the
+        # handler's handle_message() method, but internally this gives full
+        # coverage.
+        controller = UTF8Controller(self.handler)
+        controller.start()
+        self.addCleanup(controller.stop)
+
+        with SMTP(controller.hostname, controller.port) as client:
+            client.sendmail('anne@example.com', ['bart@example.com'], """\
+From: Anne Person <anne@example.com>
+To: Bart Person <bart@example.com>
+Subject: A test
+Message-ID: <ant>
+
+Testing
+""")
+        self.assertEqual(self.handled_message['subject'], 'A test')
+        self.assertEqual(self.handled_message['message-id'], '<ant>')
+        self.assertIsNotNone(self.handled_message['X-Peer'])
+        self.assertEqual(
+            self.handled_message['X-MailFrom'], 'anne@example.com')
+        self.assertEqual(self.handled_message['X-RcptTos'], 'bart@example.com')
+
+
 class TestMailbox(unittest.TestCase):
     def setUp(self):
         self.tempdir = TemporaryDirectory()