From 8de31d784b7f6957a10634540ba321509010aae4 Mon Sep 17 00:00:00 2001 From: Pierre-Olivier Mercier Date: Tue, 18 Oct 2016 01:34:46 +0200 Subject: [PATCH] Allow module function to be generators --- nemubot/treatment.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/nemubot/treatment.py b/nemubot/treatment.py index 2c1955d..884de4a 100644 --- a/nemubot/treatment.py +++ b/nemubot/treatment.py @@ -15,6 +15,7 @@ # along with this program. If not, see . import logging +import types logger = logging.getLogger("nemubot.treatment") @@ -116,10 +117,18 @@ class MessageTreater: yield r elif res is not None: - if not hasattr(res, "server") or res.server is None: - res.server = msg.server + if isinstance(res, types.GeneratorType): + for r in res: + if not hasattr(r, "server") or r.server is None: + r.server = msg.server - yield res + yield r + + else: + if not hasattr(res, "server") or res.server is None: + res.server = msg.server + + yield res hook = next(hook_gen, None)