Made AuthMembership, Group, permission customizable through Auth

emmett/tools/auth/__init__.py

@@ -1,2 +1,2 @@
 from .apis import Auth
-from .models import AuthUser
+from .models import AuthUser, AuthMembership, AuthGroup, AuthPermission

emmett/tools/auth/apis.py

@@ -27,10 +27,10 @@
 
 
 class Auth:
-    def __init__(self, app, db, user_model=None):
+    def __init__(self, app, db, user_model=None, membership_model=None, group_model=None, permission_model=None):
         self.ext = app.use_extension(AuthExtension)
         self.ext.bind_auth(self)
-        self.ext.use_database(db, user_model)
+        self.ext.use_database(db, user_model, membership_model, group_model, permission_model)
         self.ext.init_forms()
         self.pipe = AuthPipe(self)
 

emmett/tools/auth/ext.py

@@ -189,13 +189,19 @@ def _inject_pipe(self):
         if self.config.inject_pipe:
             self.app.pipeline.append(self.auth.pipe)
 
-    def use_database(self, db, user_model=None):
+    def _set_model_config(self, config_name, model):
+        if model:
+            if issubclass(model, AuthModel) or issubclass(model, TimestampedModel):
+                self.config.models[config_name] = model
+            else:
+                raise RuntimeError(f'{user_model.__name__} is an invalid user model')
+
+    def use_database(self, db, user_model=None, membership_model=None, group_model=None, permission_model=None):
         self.db = db
-        if user_model:
-            if not issubclass(user_model, AuthModel):
-                raise RuntimeError(
-                    '%s is an invalid user model' % user_model.__name__)
-            self.config.models['user'] = user_model
+        self._set_model_config("user", user_model)
+        self._set_model_config("membership", membership_model)
+        self._set_model_config("group", group_model)
+        self._set_model_config("permission", permission_model)
         self.define_models()
 
     def __set_models_labels(self):