from django import forms from django.contrib.auth.forms import UserCreationForm from django.core.exceptions import FieldDoesNotExist from django.forms import ModelForm from .models import Account, Hook, Strategy, Trade, TradingTime, User class RestrictedFormMixin: """ This mixin is used to restrict the queryset of a form to the current user. The request object is passed from the view.""" def __init__(self, *args, **kwargs): self.request = kwargs.pop("request") super().__init__(*args, **kwargs) print(self.fields) for field in self.fields: # Check it's not something like a CharField which has no queryset if not hasattr(self.fields[field], "queryset"): continue model = self.fields[field].queryset.model # Check if the model has a user field try: model._meta.get_field("user") # Add the user to the queryset filters self.fields[field].queryset = model.objects.filter( user=self.request.user ) except FieldDoesNotExist: pass class NewUserForm(UserCreationForm): email = forms.EmailField(required=True) class Meta: model = User fields = ( "username", "email", "first_name", "last_name", "password1", "password2", ) def save(self, commit=True): user = super(NewUserForm, self).save(commit=False) user.email = self.cleaned_data["email"] if commit: user.save() return user class CustomUserCreationForm(UserCreationForm): class Meta: model = User fields = "__all__" # All string/multiple choice fields class HookForm(RestrictedFormMixin, ModelForm): class Meta: model = Hook fields = ( "name", "hook", "direction", ) # All string/multiple choice fields class AccountForm(RestrictedFormMixin, ModelForm): class Meta: model = Account fields = ( "name", "exchange", "api_key", "api_secret", "sandbox", ) # Restricted mixin for account and hooks class StrategyForm(RestrictedFormMixin, ModelForm): class Meta: model = Strategy fields = ( "name", "description", "account", "trading_times", "order_type", "time_in_force", "hooks", "enabled", "take_profit_percent", "stop_loss_percent", "trailing_stop_loss_percent", "price_slippage_percent", "callback_price_deviation_percent", "trade_size_percent", ) hooks = forms.ModelMultipleChoiceField( queryset=Hook.objects.all(), widget=forms.CheckboxSelectMultiple ) trading_times = forms.ModelMultipleChoiceField( queryset=TradingTime.objects.all(), widget=forms.CheckboxSelectMultiple ) # Restricted mixin for account class TradeForm(RestrictedFormMixin, ModelForm): class Meta: model = Trade fields = ( "account", "symbol", "type", "time_in_force", "amount", "price", "stop_loss", "trailing_stop_loss", "take_profit", "direction", ) class TradingTimeForm(RestrictedFormMixin, ModelForm): class Meta: model = TradingTime fields = ( "name", "description", "start_day", "start_time", "end_day", "end_time", )