From d5c71423ba3c5905064b0f6fa29074bbe0654dcb Mon Sep 17 00:00:00 2001 From: Kevin Chung Date: Tue, 18 Aug 2020 14:59:18 -0400 Subject: [PATCH] Extract extra property into a shared function --- CTFd/forms/auth.py | 9 ++---- CTFd/forms/self.py | 26 ++++----------- CTFd/forms/users.py | 79 ++++++++++++++++++++++++++++----------------- 3 files changed, 59 insertions(+), 55 deletions(-) diff --git a/CTFd/forms/auth.py b/CTFd/forms/auth.py index 308953cf..ad3a3403 100644 --- a/CTFd/forms/auth.py +++ b/CTFd/forms/auth.py @@ -3,8 +3,8 @@ from wtforms.fields.html5 import EmailField from wtforms.validators import InputRequired from CTFd.forms import BaseForm -from CTFd.forms.users import attach_custom_user_fields from CTFd.forms.fields import SubmitField +from CTFd.forms.users import attach_custom_user_fields, build_custom_user_fields from CTFd.models import UserFields @@ -17,12 +17,7 @@ def RegistrationForm(*args, **kwargs): @property def extra(self): - fields = [] - new_fields = UserFields.query.all() - for field in new_fields: - entry = (field.name, getattr(self, f"fields[{field.id}]")) - fields.append(entry) - return fields + return build_custom_user_fields(self, include_entries=False) attach_custom_user_fields(_RegistrationForm) diff --git a/CTFd/forms/self.py b/CTFd/forms/self.py index a817d6c3..aaf9be6f 100644 --- a/CTFd/forms/self.py +++ b/CTFd/forms/self.py @@ -3,8 +3,8 @@ from wtforms import PasswordField, SelectField, StringField from wtforms.fields.html5 import DateField, URLField from CTFd.forms import BaseForm -from CTFd.forms.users import attach_custom_user_fields from CTFd.forms.fields import SubmitField +from CTFd.forms.users import attach_custom_user_fields, build_custom_user_fields from CTFd.models import FieldEntries, UserFields from CTFd.utils.countries import SELECT_COUNTRIES_LIST @@ -22,24 +22,12 @@ def SettingsForm(*args, **kwargs): @property def extra(self): - fields = [] - new_fields = UserFields.query.filter_by(editable=True).all() - user_fields = {} - - for f in FieldEntries.query.filter_by(user_id=session["id"]).all(): - user_fields[f.field_id] = f.value - - for field in new_fields: - form_field = getattr(self, f"fields[{field.id}]") - initial = user_fields.get(field.id, "") - form_field.data = initial - if form_field.render_kw: - form_field.render_kw["data-initial"] = initial - else: - form_field.render_kw = {"data-initial": initial} - entry = (field.name, form_field) - fields.append(entry) - return fields + return build_custom_user_fields( + self, + include_entries=True, + fields_kwargs={"editable": True}, + field_entries_kwargs={"user_id": session["id"]}, + ) attach_custom_user_fields(_SettingsForm, editable=True) diff --git a/CTFd/forms/users.py b/CTFd/forms/users.py index 50899e95..1013b20e 100644 --- a/CTFd/forms/users.py +++ b/CTFd/forms/users.py @@ -7,12 +7,52 @@ from CTFd.forms.fields import SubmitField from CTFd.models import FieldEntries, UserFields from CTFd.utils.countries import SELECT_COUNTRIES_LIST -def build_custom_user_fields(): - pass + +def build_custom_user_fields( + form_cls, include_entries=False, fields_kwargs=None, field_entries_kwargs=None +): + """ + Function used to reinject values back into forms for accessing by themes + """ + if fields_kwargs is None: + fields_kwargs = {} + if field_entries_kwargs is None: + field_entries_kwargs = {} + + fields = [] + new_fields = UserFields.query.filter_by(**fields_kwargs).all() + user_fields = {} + + # Only include preexisting values if asked + if include_entries is True: + for f in FieldEntries.query.filter_by(**field_entries_kwargs).all(): + user_fields[f.field_id] = f.value + + for field in new_fields: + form_field = getattr(form_cls, f"fields[{field.id}]") + + # Only include preexisting values if asked + if include_entries is True: + initial = user_fields.get(field.id, "") + form_field.data = initial + if form_field.render_kw: + form_field.render_kw["data-initial"] = initial + else: + form_field.render_kw = {"data-initial": initial} + + entry = (field.name, form_field) + fields.append(entry) + return fields def attach_custom_user_fields(form_cls, **kwargs): - new_fields = UserFields.filter_by(**kwargs).query.all() + """ + Function used to attach form fields to wtforms. + Not really a great solution but is approved by wtforms. + + https://wtforms.readthedocs.io/en/2.3.x/specific_problems/#dynamic-form-composition + """ + new_fields = UserFields.query.filter_by(**kwargs).all() for field in new_fields: validators = [] if field.required: @@ -80,24 +120,12 @@ def UserEditForm(*args, **kwargs): @property def extra(self): - fields = [] - new_fields = UserFields.query.all() - user_fields = {} - - for f in FieldEntries.query.filter_by(user_id=self.obj.id).all(): - user_fields[f.field_id] = f.value - - for field in new_fields: - form_field = getattr(self, f"fields[{field.id}]") - initial = user_fields.get(field.id, "") - form_field.data = initial - if form_field.render_kw: - form_field.render_kw["initial"] = initial - else: - form_field.render_kw = {"data-initial": initial} - entry = (field.name, form_field) - fields.append(entry) - return fields + return build_custom_user_fields( + self, + include_entries=True, + fields_kwargs=None, + field_entries_kwargs={"user_id": self.obj.id}, + ) def __init__(self, *args, **kwargs): """ @@ -119,14 +147,7 @@ def UserCreateForm(*args, **kwargs): @property def extra(self): - fields = [] - new_fields = UserFields.query.all() - - for field in new_fields: - form_field = getattr(self, f"fields[{field.id}]") - entry = (field.name, form_field) - fields.append(entry) - return fields + return build_custom_user_fields(self, include_entries=False) attach_custom_user_fields(_UserCreateForm)