fixed test related to getting the correct year from a Month()

5 files changed, 43 insertions, 30 deletions

diff --git a/datamaps/api/api.py b/datamaps/api/api.py
index 4073fbc..072d95f 100644
--- a/datamaps/api/api.py
+++ b/datamaps/api/api.py
@@ -15,7 +15,7 @@ def project_data_from_master_api(master_file: str, quarter: int, year: int):
     return m
 
 
-def project_data_from_master_month_api(master_file: str, month: int, year: int):
+def project_data_from_master_month_api(master_file: str, month: int, year: int) -> Master:
     """Create a Master object directly without the need to explicitly pass
     a Month object.
 
diff --git a/datamaps/core/temporal.py b/datamaps/core/temporal.py
index f97edfa..a927b23 100644
--- a/datamaps/core/temporal.py
+++ b/datamaps/core/temporal.py
@@ -1,6 +1,22 @@
 import calendar
 import datetime
 import itertools
+from typing import List
+
+MONTHS = [
+    "January",
+    "February",
+    "March",
+    "April",
+    "May",
+    "June",
+    "July",
+    "August",
+    "September",
+    "October",
+    "November",
+    "December",
+]
 
 
 class Month:
@@ -22,40 +38,27 @@ class Month:
     }
 
     def __init__(self, month: int, year: int):
-        self._month_int = month
+        self.month_int = month
         self.year = year
 
     @property
     def start_date(self):
-        return datetime.date(self.year, self._month_int, 1)
+        return datetime.date(self.year, self.month_int, 1)
 
     @property
     def end_date(self):
-        if self._month_int == 2 and calendar.isleap(self.year):
+        if self.month_int == 2 and calendar.isleap(self.year):
             return datetime.date(
-                self.year, self._month_int, Month._end_ints[self._month_int] + 1
+                self.year, self.month_int, Month._end_ints[self.month_int] + 1
             )
         else:
             return datetime.date(
-                self.year, self._month_int, Month._end_ints[self._month_int]
+                self.year, self.month_int, Month._end_ints[self.month_int]
             )
 
     @property
     def name(self):
-        return [
-            "January",
-            "February",
-            "March",
-            "April",
-            "May",
-            "June",
-            "July",
-            "August",
-            "September",
-            "October",
-            "November",
-            "December",
-        ][self._month_int - 1]
+        return MONTHS[self.month_int - 1]
 
     def __repr__(self):
         return f"Month({self.name})"
@@ -86,7 +89,7 @@ class FinancialYear:
         self.end_date = self.q4.end_date
 
     @property
-    def q1(self) -> datetime.date:
+    def q1(self):
         """Quarter 1 as a :py:class:`datetime.date` object"""
         return self._q1
 
@@ -156,7 +159,7 @@ class Quarter:
 
         self.start_date = self._start_date(self.quarter, self.year)
         self.end_date = self._end_date(self.quarter, self.year)
-        self.months = []
+        self.months: List[Month] = []
         self._populate_months()
 
     def _populate_months(self):
diff --git a/datamaps/plugins/dft/master.py b/datamaps/plugins/dft/master.py
index ac81bd9..324b63d 100644
--- a/datamaps/plugins/dft/master.py
+++ b/datamaps/plugins/dft/master.py
@@ -1,14 +1,13 @@
-import re
 import datetime
 import logging
+import re
 import unicodedata
 from pathlib import Path
-from typing import List, Tuple, Iterable, Optional, Any
+from typing import Any, Iterable, List, Optional, Tuple
 
+from datamaps.core.temporal import Quarter
 from datamaps.plugins.dft.portfolio import project_data_from_master
 from datamaps.process.cleansers import DATE_REGEX_4
-from datamaps.core.temporal import Quarter
-
 from openpyxl import load_workbook
 
 logger = logging.getLogger("bcompiler.utils")
@@ -157,9 +156,14 @@ class Master:
         self._quarter = quarter
         self._declared_month = declared_month
         self.path = path
+        if self._declared_month:
+            idxs = [x.month_int for x in self._quarter.months]
+            m_idx = idxs.index(self._declared_month)
+            self.year = self._quarter.months[m_idx].year
+        else:
+            self.year = self._quarter.year
         self._data = project_data_from_master(self.path)
         self._project_titles = [item for item in self.data.keys()]
-        self.year = self._quarter.year
 
     def __getitem__(self, project_name):
         return ProjectData(self._data[project_name])
@@ -217,9 +221,7 @@ class Master:
             12: "December",
         }
         return [
-            m
-            for m in self.quarter.months
-            if m.name == months[self._declared_month]
+            m for m in self.quarter.months if m.name == months[self._declared_month]
         ][0]
 
     @property
diff --git a/datamaps/tests/test_api.py b/datamaps/tests/test_api.py
index 4c73931..295f1fb 100644
--- a/datamaps/tests/test_api.py
+++ b/datamaps/tests/test_api.py
@@ -18,6 +18,7 @@ def test_get_project_data_using_month(master):
     m2 = project_data_from_master_month(master, 8, 2021)
     m3 = project_data_from_master_month(master, 9, 2021)
     m4 = project_data_from_master_month(master, 10, 2021)
+    m5 = project_data_from_master_month(master, 2, 2021)  # this is q4
     assert m["Chutney Bridge.xlsm"]["Project/Programme Name"] == "Chutney Bridge Ltd"
     assert isinstance(m.month, Month)
     assert m.month.name == "July"
@@ -32,3 +33,9 @@ def test_get_project_data_using_month(master):
     assert m2.quarter.end_date == datetime.date(2021, 9, 30)
     assert m3.quarter.end_date == datetime.date(2021, 9, 30)
     assert m4.quarter.end_date == datetime.date(2021, 12, 31)
+
+    # year should be different if using this func
+    assert m.year == 2021
+    assert m2.year == 2021
+    assert m3.year == 2021
+    assert m5.year == 2021
diff --git a/datamaps/tests/test_month.py b/datamaps/tests/test_month.py
index a533cd0..de53545 100644
--- a/datamaps/tests/test_month.py
+++ b/datamaps/tests/test_month.py
@@ -30,6 +30,7 @@ def test_quarter_objects_have_months():
 def test_month():
     m1 = Month(1, 2021)
     assert m1.name == "January"
+    assert m1.year == 2021
     m2 = Month(9, 2021)
     assert m2.name == "September"
     assert m2.start_date == datetime.date(2021, 9, 1)