diff --git a/helpers/answer.py b/helpers/answer.py index 5037f66..6998372 100644 --- a/helpers/answer.py +++ b/helpers/answer.py @@ -20,9 +20,6 @@ class BaseAnswer: user = User(chat_id=message.chat.id) mongo.users_collection.insert_one(UserSchema().dump(user)) else: - # что-то со временем в маршмеллоу происходит, не знаю как решить - if 'next_daily_notify_time' in user: - user['next_daily_notify_time'] = str(user['next_daily_notify_time']) user = UserSchema().load(user) attr = getattr(self, "handle_state_" + user.state, None) if attr is None: diff --git a/helpers/models.py b/helpers/models.py index f0ecc94..6c3391e 100644 --- a/helpers/models.py +++ b/helpers/models.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from typing import Optional import marshmallow_dataclass -from marshmallow import EXCLUDE +from marshmallow import EXCLUDE, post_load @dataclass @@ -27,4 +27,13 @@ class Lesson: -UserSchema = marshmallow_dataclass.class_schema(User) +USchema = marshmallow_dataclass.class_schema(User) + + +class UserSchema(USchema): + + @post_load + def stringify_time(self, in_data, **kwargs): + if 'next_daily_notify_time' in in_data: + in_data['next_daily_notify_time'] = str(in_data['next_daily_notify_time']) + return in_data