From a27b87613b6395bbcb7be1eac510cab21c19e452 Mon Sep 17 00:00:00 2001
From: h-janes <dev@hjanes.com>
Date: Mon, 28 Nov 2022 20:46:38 +0000
Subject: [PATCH] lastmod arg supports datetime object

---
 flask_sitemapper/sitemapper.py | 13 ++++++--
 tests/test_datetime.py         | 58 ++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_datetime.py

diff --git a/flask_sitemapper/sitemapper.py b/flask_sitemapper/sitemapper.py
index 0c2c04b..52536e8 100644
--- a/flask_sitemapper/sitemapper.py
+++ b/flask_sitemapper/sitemapper.py
@@ -1,5 +1,6 @@
 """Provides the `URL` and `Sitemapper` classes"""
 
+from datetime import datetime
 from functools import wraps
 from typing import Callable, Union
 
@@ -17,7 +18,7 @@ def __init__(
         self,
         endpoint,
         scheme: str,
-        lastmod: str = None,
+        lastmod: Union[str, datetime] = None,
         changefreq: str = None,
         priority: Union[str, int, float] = None,
     ) -> None:
@@ -27,6 +28,9 @@ def __init__(
         self.changefreq = changefreq
         self.priority = priority
 
+        if isinstance(self.lastmod, datetime):
+            self.lastmod = self.lastmod.strftime("%Y-%m-%dT%H:%M:%S")
+
     @property
     def loc(self) -> str:
         """Finds the URL from the endpoint name. Must be called within a request context"""
@@ -77,7 +81,10 @@ def init_app(self, app: Flask) -> None:
         self.deferred_functions.clear()
 
     def include(
-        self, lastmod: str = None, changefreq: str = None, priority: Union[str, int, float] = None
+        self,
+        lastmod: Union[str, datetime] = None,
+        changefreq: str = None,
+        priority: Union[str, int, float] = None,
     ) -> Callable:
         """A decorator for view functions to add their URL to the sitemap"""
 
@@ -109,7 +116,7 @@ def __get_endpoint_name(self, func: Callable) -> str:
     def add_endpoint(
         self,
         view_func: Union[Callable, str],
-        lastmod: str = None,
+        lastmod: Union[str, datetime] = None,
         changefreq: str = None,
         priority: Union[str, int, float] = None,
     ) -> None:
diff --git a/tests/test_datetime.py b/tests/test_datetime.py
new file mode 100644
index 0000000..4e42baf
--- /dev/null
+++ b/tests/test_datetime.py
@@ -0,0 +1,58 @@
+from datetime import datetime
+
+import flask
+
+from flask_sitemapper import Sitemapper
+
+# ----------------- TEST APP -----------------
+
+sitemapper = Sitemapper()
+app = flask.Flask(__name__)
+sitemapper.init_app(app)
+
+
+@sitemapper.include(lastmod=datetime(2022, 11, 28))
+@app.route("/")
+def r_home():
+    return "<h1>Home</h1>"
+
+
+@app.route("/sitemap.xml")
+def r_sitemap():
+    return sitemapper.generate()
+
+
+# ----------------- END TEST APP -----------------
+
+
+def test_running():
+    with app.test_client() as test_client:
+        response = test_client.get("/")
+        assert response.text == "<h1>Home</h1>"
+
+
+def test_status_code():
+    with app.test_client() as test_client:
+        response = test_client.get("/sitemap.xml")
+        assert response.status_code == 200
+
+
+def test_mimetype():
+    with app.test_client() as test_client:
+        response = test_client.get("/sitemap.xml")
+        assert response.mimetype == "application/xml"
+
+
+def test_xml():
+    with app.test_client() as test_client:
+        response = test_client.get("/sitemap.xml")
+        assert (
+            response.text
+            == """<?xml version="1.0" encoding="utf-8"?>
+<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
+  <url>
+    <loc>https://localhost/</loc>
+    <lastmod>2022-11-28T00:00:00</lastmod>
+  </url>
+</urlset>"""
+        )