From 6ebabbc51ca92a1af334f5e8983e56bdb8b9cc7f Mon Sep 17 00:00:00 2001
From: Scoder12 <34356756+Scoder12@users.noreply.github.com>
Date: Tue, 28 Jul 2020 16:17:14 -0700
Subject: [PATCH] Feat: Add needs_sigin decorator

---
 src/replit/maqpy/utils.py | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/src/replit/maqpy/utils.py b/src/replit/maqpy/utils.py
index 7935a39..07b710e 100644
--- a/src/replit/maqpy/utils.py
+++ b/src/replit/maqpy/utils.py
@@ -1,4 +1,9 @@
 """Utitilities to make development easier."""
+from typing import Any, Callable
+from functools import wraps
+
+import flask
+
 from .html import Page
 
 
@@ -18,3 +23,28 @@ def signin(title: str = "Please Sign In") -> Page:
         Page: The sign-in page.
     """
     return Page(title=title, body=sign_in_snippet)
+
+
+def needs_signin(func: Callable = None, loginhtml: str = sign_in_snippet) -> Callable:
+    """A decorator that enforces that the user is signed in before accessing the page.
+
+    Args:
+        func (Callable): The function passed in if used as a decorator. Defaults to None.
+        loginhtml (str): The HTML to show when the user needs to sign in. Defaults to sign_in_snippet.
+
+    Returns:
+        Callable: The new handler.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def handler(*args: Any, **kwargs: Any) -> flask.Response:
+            if flask.request.signed_in:
+                return func(*args, **kwargs)
+            else:
+                return loginhtml
+
+    if func is not None:  # called with no options @needs_signin
+        return decorator(func)
+    else:  # called with options, eg @needs_signin(loginhtml='...')
+        return decorator