fisk/core/forms.py

145 lines
3.7 KiB
Python

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",
)