1 files changed, 27 insertions, 4 deletions
diff --git a/apps/notes/forms.py b/apps/notes/forms.py
index 5ef9f84..4dc79d2 100644
--- a/apps/notes/forms.py
+++ b/apps/notes/forms.py
@@ -1,10 +1,19 @@
 from django import forms
 from django.utils.translation import ugettext_lazy as _
+from django.core.exceptions import NON_FIELD_ERRORS
+
+from utils.widgets import RelatedFieldWidgetCanAdd
 
 from .models import Note, Notebook
 
 
-class NoteForm(forms.ModelForm):
+class BaseNoteForm(forms.ModelForm):
+    def __init__(self, *args, **kwargs):
+        self.user = kwargs.pop("user", None)
+        super(BaseNoteForm, self).__init__(*args, **kwargs)
+
+
+class NoteForm(BaseNoteForm):
     class Meta:
         model = Note
         fields = ['title', 'body_text', 'body_html', 'body_qjson', 'notebook', 'url', 'tags']
@@ -13,14 +22,28 @@ class NoteForm(forms.ModelForm):
         }
 
     def __init__(self, *args, **kwargs):
-        self.user = kwargs.pop("user", None)
+        user = kwargs.pop("user", None)
         super(NoteForm, self).__init__(*args, **kwargs)
+        self.fields['notebook'].widget = RelatedFieldWidgetCanAdd(Notebook, related_url="notebooks:list")
+        self.fields['notebook'].queryset = Notebook.objects.filter(owner__username=user)
 
 
-class NotebookForm(NoteForm):
+class NotebookForm(BaseNoteForm):
     class Meta:
         model = Notebook
-        fields = ['name', 'color_rgb']
+        fields = ['owner', 'name', 'color_rgb']
+        widgets = {'owner': forms.HiddenInput()}
         labels = {
             "name": _("Notebook Name"),
+            "color_rgb": _("Notebook Color"),
+        }
+        error_messages = {
+            NON_FIELD_ERRORS: {
+                'unique_together': "You already have a notebook by that name, please choose a different name",
+            }
         }
+
+    def __init__(self, *args, **kwargs):
+        user = kwargs.pop("user", None)
+        super(NotebookForm, self).__init__(*args, **kwargs)
+        self.fields['owner'].initial = user