Allow module function to be generators

This commit is contained in:
nemunaire 2016-10-18 01:34:46 +02:00 committed by nemunaire
parent 2e5834a89d
commit 6d2f90fe77

View File

@ -15,6 +15,7 @@
# along with this program. If not, see <http://www.gnu.org/licenses/>.
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)