From 1c76c8d91169401cbc752a4c614940d4df1d1bb4 Mon Sep 17 00:00:00 2001
From: profitroll <vozhd.kk@gmail.com>
Date: Thu, 30 May 2024 23:22:29 +0200
Subject: [PATCH] Added automatic timezone update on location change; Update
 methods now return changed values

---
 classes/pyrouser.py | 68 +++++++++++++++++++++++++++++++++++----------
 1 file changed, 54 insertions(+), 14 deletions(-)

diff --git a/classes/pyrouser.py b/classes/pyrouser.py
index 9353e6e..2a91a8a 100644
--- a/classes/pyrouser.py
+++ b/classes/pyrouser.py
@@ -1,8 +1,9 @@
 import logging
 from dataclasses import dataclass
-from datetime import datetime, timedelta, timezone
-from typing import Any, Union
+from datetime import datetime, timedelta
+from typing import Any, Mapping, Tuple, Union
 
+import pytz
 from bson import ObjectId
 
 from classes.location import Location
@@ -45,7 +46,7 @@ class PyroUser:
         offset: int = 1,
         time_hour: int = 16,
         time_minute: int = 0,
-    ):
+    ) -> "PyroUser":
         db_entry = await col_users.find_one({"id": id})
 
         if db_entry is None:
@@ -73,7 +74,7 @@ class PyroUser:
         return cls(**db_entry)
 
     @classmethod
-    async def from_dict(cls, **kwargs):
+    async def from_dict(cls, **kwargs) -> "PyroUser":
         if "location" in kwargs:
             try:
                 kwargs["location"] = await Location.get(kwargs["location"])  # type: ignore
@@ -81,46 +82,85 @@ class PyroUser:
                 kwargs["location"] = None  # type: ignore
         return cls(**kwargs)
 
-    async def update_locale(self, locale: Union[str, None]) -> None:
+    async def update_locale(self, locale: Union[str, None]) -> Union[str, None]:
         """Change user's locale stored in the database.
 
         ### Args:
             * locale (`Union[str, None]`): New locale to be set.
         """
+
         logger.debug("%s's locale has been set to %s", self.id, locale)
+
         await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})
+
         self.locale = locale
 
-    async def update_state(self, enabled: bool = False) -> None:
+        return self.locale
+
+    async def update_state(self, enabled: bool = False) -> bool:
         logger.debug("%s's state has been set to %s", self.id, enabled)
+
         await col_users.update_one({"_id": self._id}, {"$set": {"enabled": enabled}})
+
         self.enabled = enabled
 
-    async def update_location(self, location_id: int = 0) -> None:
+        return self.enabled
+
+    async def update_location(self, location_id: int = 0) -> Location:
         logger.debug("%s's location has been set to %s", self.id, location_id)
+
         await col_users.update_one(
             {"_id": self._id}, {"$set": {"location": location_id}}
         )
-        self.location = await Location.get(location_id)
 
-    async def update_offset(self, offset: int = 1) -> None:
+        location = await Location.get(location_id)
+
+        # Execute if timezones of old and new locations are different
+        if self.location and (self.location.timezone.zone != location.timezone.zone):
+            # Get UTC time for selected reminder time
+            now_utc = datetime.now(pytz.utc).replace(
+                hour=self.time_hour, minute=self.time_minute, second=0, microsecond=0
+            )
+
+            # Get the time for the reminder time of old and new location
+            local_old = now_utc.astimezone(self.location.timezone)
+            local_new = (
+                location.timezone.localize(local_old.replace(tzinfo=None))
+            ).astimezone(pytz.utc)
+
+            # Update the time to match the new timezone
+            await self.update_time(hour=local_new.hour, minute=local_new.minute)
+
+        self.location = location
+
+        return self.location
+
+    async def update_offset(self, offset: int = 1) -> int:
         logger.debug("%s's offset has been set to %s", self.id, offset)
+
         await col_users.update_one({"_id": self._id}, {"$set": {"offset": offset}})
+
         self.offset = offset
 
-    async def update_time(self, hour: int = 16, minute: int = 0) -> None:
+        return offset
+
+    async def update_time(self, hour: int = 16, minute: int = 0) -> Tuple[int, int]:
         logger.debug("%s's time has been set to %s h. %s m.", self.id, hour, minute)
+
         await col_users.update_one(
             {"_id": self._id}, {"$set": {"time_hour": hour, "time_minute": minute}}
         )
+
         self.time_hour = hour
         self.time_minute = minute
 
+        return self.time_hour, self.time_minute
+
     async def delete(self) -> None:
         logger.debug("%s's data has been deleted", self.id)
         await col_users.delete_one({"_id": self._id})
 
-    async def checkout(self) -> Any:
+    async def checkout(self) -> Mapping[str, Any]:
         logger.debug("%s's data has been checked out", self.id)
         db_entry = await col_users.find_one({"_id": self._id})
 
@@ -155,7 +195,7 @@ class PyroUser:
             logger.warning("Location %s does not have a timezone set", self.location.id)
 
         return (
-            datetime.now(self.location.timezone or timezone.utc) + timedelta(days=1)
+            datetime.now(self.location.timezone or pytz.utc) + timedelta(days=1)
         ).replace(hour=0, minute=0, second=0, microsecond=0)
 
     def get_reminder_time(self) -> datetime:
@@ -180,12 +220,12 @@ class PyroUser:
             logger.warning("Location %s does not have a timezone set", self.location.id)
 
         return (
-            datetime.now(timezone.utc)
+            datetime.now(pytz.utc)
             .replace(
                 hour=self.time_hour,
                 minute=self.time_minute,
                 second=0,
                 microsecond=0,
             )
-            .astimezone(self.location.timezone or timezone.utc)
+            .astimezone(self.location.timezone or pytz.utc)
         )