add registration

1 files changed, 16 insertions, 1 deletions

diff --git a/app/models.py b/app/models.py
index 9d11ffb..b607748 100644
--- a/app/models.py
+++ b/app/models.py
@@ -1,11 +1,15 @@
 
 import datetime
 
+from flask_login import UserMixin
+
+from werkzeug.security import generate_password_hash, check_password_hash
 
 from app import db
+from app import login
 
 
-class User(db.Model):
+class User(UserMixin, db.Model):
     id = db.Column(db.Integer, primary_key=True)
     username = db.Column(db.String(80), unique=True, nullable=False)
     email = db.Column(db.String(), unique=True, nullable=False)
@@ -17,6 +21,12 @@ class User(db.Model):
     def __repr__(self):
         return "<User %r>" % self.username
 
+    def set_password(self, password):
+        self.password_hash = generate_password_hash(password)
+
+    def check_password(self, password):
+        return check_password_hash(self.password_hash, password)
+
 
 class Dataset(db.Model):
     id = db.Column(db.Integer, primary_key=True)
@@ -57,3 +67,8 @@ class Annotation(db.Model):
 
     def __repr__(self):
         return "<Annotation %r>" % self.id
+
+
+@login.user_loader
+def load_user(_id):
+    return User.query.get(int(_id))