Merge pull request #7 from iMHouse/master

Changes to calculator.py
swimlane · Sep 12, 2023 · 61dc36c · 61dc36c
2 parents 4909b15 + fec86cd
commit 61dc36c
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/sla_calculator/calculator.py b/sla_calculator/calculator.py
@@ -7,23 +7,19 @@ def __init__(self):
         pass
 
     def check_working_days(self, start_time, country_holidays, open_hour):
-        while start_time in country_holidays:
+        while (start_time.day_of_week in (pendulum.SUNDAY, pendulum.SATURDAY)) or (pendulum.date(start_time.year, start_time.month, start_time.day) in country_holidays):
             start_time = start_time.add(days=1)
             start_time = pendulum.datetime(start_time.year, start_time.month, start_time.day, open_hour, 0, 0,
-                                           tz=pendulum.local_timezone())
-        while start_time.day_of_week > 4:  # 0=Monday, 6=Sunday
-            start_time = start_time.add(days=1)
-            start_time = pendulum.datetime(start_time.year, start_time.month, start_time.day, open_hour, 0, 0,
-                                           tz=pendulum.local_timezone())
+                                            tz=pendulum.local_timezone())
         return start_time
 
     def calculate(self, start_time, open_hour, close_hour, country_name='US', sla_in_hours=4, province=None,
                   state=None):
         sla_time = None
         sla_in_minutes = sla_in_hours * 60
-        start_time = start_time if isinstance(start_time, pendulum.DateTime) else pendulum.parse(start_time)
+        start_time = start_time if isinstance(start_time, pendulum.DateTime) else pendulum.parse(start_time, tz = pendulum.local_timezone())
 
-        country_holidays = pyholidays.CountryHoliday(country_name, prov=province, state=state)
+        country_holidays = list(pyholidays.CountryHoliday(country_name, years=[start_time.year], prov=province, state=state).keys())
         start_time = self.check_working_days(start_time, country_holidays, open_hour)
 
         open_time = pendulum.datetime(start_time.year, start_time.month, start_time.day, open_hour,
@@ -34,6 +30,9 @@ def calculate(self, start_time, open_hour, close_hour, country_name='US', sla_in
             start_time = open_time
         elif start_time > close_time:
             start_time = open_time.add(days=1)
+
+            start_time = self.check_working_days(start_time, country_holidays, open_hour)
+
             open_time = pendulum.datetime(start_time.year, start_time.month, start_time.day, open_hour, 0, 0,
                                           tz=pendulum.local_timezone())
             close_time = pendulum.datetime(start_time.year, start_time.month, start_time.day, close_hour, 0, 0,