Repository: BraeWebb/allocate

Issue:
Availabilities Format
The only accepted format for availabilities at the moment is the output of a doodle poll. This format is kind of gross and it would be better to have a standard availability format with the ability to accept the doodle format being optional.

With standard availability format:
```bash
allocate tutors.csv sessions.csv availability.csv
```

With a doodle availability format (functions the same as the current system):
```bash
allocate --doodle tutors.csv sessions.csv availability.csv
```

Comments:

Patch:
diff --git a/README.md b/README.md
index 7d8a697..551f73a 100644
--- a/README.md
+++ b/README.md
@@ -74,17 +74,17 @@ The following tables represent the 3 example CSV files that come with the projec
 | P05 | 2                 | 2                 | Thu | 9          | 2        |
 
 ### availabilities.csv
-| "Poll ""Sample Doodle Availability""" |   |   |   |  |  |  |  |  |  |
-|---|---|---|---|---|---|---|---|---|---|
-| https://doodle\.com/poll/v7scx5pzwgeu3m2y |   |   |   |  |  |  |  |  |  |
-|   |   |   |   |  |  |  |  |  |  |
-|   | February 2020 |   |   |  |  |  |  |  |  |
-|   | Tue 4 |   |   | Wed 5 |  |  | Thu 6 | Fri 7 |  |
-|   | 11:00 AM – 1:00 PM | 1:00 PM – 3:00 PM | 3:00 PM – 5:00 PM | 8:00 AM – 9:00 AM | 9:00 AM – 10:00 AM | 12:00 PM – 2:00 PM | 9:00 AM – 11:00 AM | 11:00 AM – 12:00 PM | 12:00 PM – 1:00 PM |
-| Brae | OK |   | OK |  | OK | OK | OK |  | OK |
-| Henry | OK | OK | OK | OK |  |  | OK |  |  |
-| Emily |   | OK |   |  | OK | OK |  | OK | OK |
-| Count | 2 | 2 | 1 | 2 | 2 | 2 | 2 | 2 | 2 |
+The first row is the day of the week of a session, second row represents the
+time of day in 24 hours and the third row is the duration of the session.
+Anything but a 1 in the availability cell of each tutor is treated as unavailable.
+
+|       | Tue | Tue | Tue | Wed | Wed | Wed | Thu | Fri | Fri |
+|-------|-----|-----|-----|-----|-----|-----|-----|-----|-----|
+|       | 11  | 13  | 15  | 8   | 9   | 12  | 9   | 11  | 12  |
+|       | 2   | 2   | 2   | 1   | 1   | 2   | 2   | 1   | 1   |
+| Brae  | 1   |     | 1   |     | 1   | 1   | 1   |     | 1   |
+| Henry | 1   | 1   | 1   | 1   |     |     | 1   |     |     |
+| Emily |     | 1   |     |     | 1   | 1   |     | 1   | 1   |
 
 With the above files available in the current working directory, execution of the following command will produce the allocation below
 ```bash
diff --git a/allocate/allocation.py b/allocate/allocation.py
index cf9989f..c6f72a4 100644
--- a/allocate/allocation.py
+++ b/allocate/allocation.py
@@ -7,7 +7,7 @@ from typing import Dict, List
 from allocate.solver import validate_availability, Engine
 from allocate.model import Tutor, Session
 from allocate.csvalidator import CSVModel
-from allocate.doodle import parse_doodle_to_stub, parse_doodle_hack
+from allocate.availability import Availability
 
 
 def solution_to_csv(solution: Dict[str, List[str]], output):
@@ -20,10 +20,14 @@ def solution_to_csv(solution: Dict[str, List[str]], output):
         writer.writerow([name, *sessions])
 
 
-def stub_files(tutors: str, sessions: str, availability: str):
+def stub_files(tutors: str, sessions: str, availability: str,
+               doodle: bool = False):
     """Write out stub files for tutors and sessions based on the
     given availability file."""
-    tutors_names, sessions_details = parse_doodle_to_stub(availability)
+    if doodle:
+        availability_data = Availability.from_doodle(availability)
+    else:
+        availability_data = Availability.from_csv(availability)
 
     with open(tutors, 'x') as file:
         writer = csv.writer(file)
@@ -32,7 +36,7 @@ def stub_files(tutors: str, sessions: str, availability: str):
         writer.writerow(Tutor.__annotations__.keys())
         columns = len(Tutor.__annotations__)
         # write tutor names
-        for tutor in sorted(tutors_names):
+        for tutor in sorted(availability_data.tutors):
             writer.writerow([tutor] + ["" for _ in range(columns - 1)])
 
     with open(sessions, 'x') as file:
@@ -41,24 +45,28 @@ def stub_files(tutors: str, sessions: str, availability: str):
         # write header
         writer.writerow(Session.__annotations__.keys())
         # write session names
-        for session in sessions_details:
+        for session in availability_data.sessions:
             writer.writerow(["", session.day, session.start, session.duration, "", ""])
 
 
 def run_allocation(tutors: str, sessions: str, availability: str,
-                   json: bool = False):
+                   doodle: bool = False, json: bool = False):
     tutor_model = CSVModel(Tutor)
     tutor_model.load(tutors, allow_defaults=True)
 
     session_model = CSVModel(Session)
     session_model.load(sessions, allow_defaults=True)
 
-    availability_data = parse_doodle_hack(availability, tutor_model, session_model)
+    if doodle:
+        availability_data = Availability.from_doodle(availability)
+    else:
+        availability_data = Availability.from_csv(availability)
+    availability_matrix = availability_data.to_matrix(tutor_model, session_model)
 
-    for message in validate_availability(availability_data):
+    for message in validate_availability(availability_matrix):
         print(message)
 
-    engine = Engine(tutor_model, session_model, availability_data)
+    engine = Engine(tutor_model, session_model, availability_matrix)
     solution = engine.solve()
 
     if solution is None:
@@ -83,6 +91,8 @@ def main():
     parser.add_argument('availability', type=str,
                         help='CSV file of tutors availabilities to sessions')
 
+    parser.add_argument('--doodle', action='store_true',
+                        help='Parse the input availability table as a doodle export')
     parser.add_argument('--stub', action='store_true',
                         help='Write to the tutor and session files a stub table generated by the allocations file')
     parser.add_argument('--json', action="store_true",
@@ -91,9 +101,11 @@ def main():
     args = parser.parse_args()
 
     if args.stub:
-        stub_files(args.tutors, args.sessions, args.availability)
+        stub_files(args.tutors, args.sessions, args.availability,
+                   doodle=args.doodle)
     else:
-        run_allocation(args.tutors, args.sessions, args.availability, json=args.json)
+        run_allocation(args.tutors, args.sessions, args.availability,
+                       doodle=args.doodle, json=args.json)
 
 
 if __name__ == '__main__':
diff --git a/allocate/availability.py b/allocate/availability.py
new file mode 100644
index 0000000..90eac30
--- /dev/null
+++ b/allocate/availability.py
@@ -0,0 +1,135 @@
+"""Handles saving and loading of an availability file"""
+import io
+import csv
+from collections import defaultdict
+from typing import Dict, Iterable, Tuple, List, Mapping
+
+from allocate.model import Tutor, Session, TimeSlot
+from allocate.doodle import parse_doodle_to_stub, parse_doodle
+
+
+class Availability:
+    """Availability represents the times a tutor indicates they are available."""
+
+    def __init__(self):
+        self._tutors_to_times: Mapping[str, List[TimeSlot]] = {}
+        self._all_slots: Iterable[TimeSlot] = []
+
+    @classmethod
+    def from_doodle(cls, doodle_file: str):
+        """Load an availability from a doodle CSV file"""
+        instance = cls()
+        instance._tutors_to_times = parse_doodle(doodle_file)
+        instance._all_slots = parse_doodle_to_stub(doodle_file)[1]
+        return instance
+
+    @classmethod
+    def from_csv(cls, filename: str):
+        """Load an availability from a regular availability CSV file"""
+        instance = cls()
+
+        with open(filename, 'r') as file:
+            reader = iter(csv.reader(file))
+
+            day_row = next(reader)[1:]
+            time_row = map(int, next(reader)[1:])
+            duration_row = map(int, next(reader)[1:])
+
+            days = [TimeSlot(*slot) for slot in zip(day_row, time_row, duration_row)]
+
+            availabilities: Mapping[str, List[TimeSlot]] = defaultdict(list)
+            for row in reader:
+                name = row[0]
+
+                # add every availability timeslot
+                for column, status in enumerate(row):
+                    if status == "1":
+                        availabilities[name].append(days[column - 1])
+
+        instance._tutors_to_times = availabilities
+        instance._all_slots = days
+
+        return instance
+
+    def get_available_slots(self, tutor: str) -> Iterable[TimeSlot]:
+        """Get the time slots a tutor is available"""
+        return self._tutors_to_times[tutor]
+
+    def is_available(self, tutor: str, session: TimeSlot) -> bool:
+        """Determine if a tutor is available at the given timeslot"""
+        return session in self.get_available_slots(tutor)
+
+    @property
+    def tutors(self) -> Iterable[str]:
+        """All the tutors in this availability"""
+        return self._tutors_to_times.keys()
+
+    @property
+    def sessions(self) -> Iterable[TimeSlot]:
+        """All the sessions in this availability"""
+        return self._all_slots
+
+    def write(self, filename: str):
+        """Write the availability information to a CSV file"""
+        with open(filename, 'w') as file:
+            file.write(self.to_csv())
+
+    def to_matrix(self, tutors: Iterable[Tutor], sessions: Iterable[Session])\
+            -> Dict[Tuple[Tutor, Session], bool]:
+        matrix = {}
+        session_map = {TimeSlot(session.day.value, session.start_time,
+                                session.duration): session for session in sessions}
+
+        for slot, session in session_map.items():
+            for tutor in tutors:
+                matrix[(tutor, session)] = self.is_available(tutor.name, slot)
+
+        return matrix
+
+    def to_csv(self) -> str:
+        """Convert the availability information to a CSV file format"""
+        # calculate the index of each time slow
+        columns = {}
+        current_column = 1
+        for session in self._all_slots:
+            columns[session] = current_column
+            current_column += 1
+
+        output = io.StringIO()
+        writer = csv.writer(output)
+
+        # write the columns for each time slot
+        data = [(session.day, session.start, session.duration)
+                for session in self._all_slots]
+        for slot_row in zip(*data):
+            writer.writerow(("",) + slot_row)
+
+        # write in availabilities
+        blank_row = [""] + ["" for _ in self._all_slots]
+        for name, times in self._tutors_to_times.items():
+            row = blank_row[:]
+            row[0] = name
+
+            for available_slot in times:
+                slot_index = columns[available_slot]
+                row[slot_index] = "1"
+
+            writer.writerow(row)
+
+        return output.getvalue()
+
+
+def main():
+    availability = Availability.from_doodle("sample_doodle.csv")
+    availability.write("tmp.csv")
+
+    written = Availability.from_csv("tmp.csv")
+
+    print(written.to_csv() == availability.to_csv())
+
+    import os
+    os.remove("tmp.csv")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/allocate/doodle.py b/allocate/doodle.py
index 885ffeb..850fe63 100644
--- a/allocate/doodle.py
+++ b/allocate/doodle.py
@@ -4,20 +4,11 @@ that other programs are able to parse with relative ease.
 """
 
 import csv
-from typing import Dict, Iterable, Tuple, List
-from dataclasses import dataclass
+from typing import Mapping, Iterable, Tuple, List
 from collections import defaultdict
 from itertools import islice
 
-from allocate.model import Tutor, Session
-
-
-@dataclass(eq=True, frozen=True)
-class TimeSlot:
-    """Data class used to represent a class that occupied a time slot"""
-    day: str
-    start: int
-    duration: int
+from allocate.model import TimeSlot
 
 
 def _hour_from_time(time: str) -> int:
@@ -75,9 +66,9 @@ def _assign_columns_timeslots(days: Iterable[str], times: Iterable[str]):
     return columns
 
 
-# TODO: Abstract the 3 functions below
+# TODO: Abstract the 2 functions below
 
-def parse_doodle(filename: str) -> Dict[str, List[TimeSlot]]:
+def parse_doodle(filename: str) -> Mapping[str, List[TimeSlot]]:
     """Parse a Doodle CSV to create a dictionary that maps a tutor name
     to a list of all their available time slots.
     """
@@ -91,7 +82,7 @@ def parse_doodle(filename: str) -> Dict[str, List[TimeSlot]]:
         time_row = next(reader)
         days = _assign_columns_timeslots(day_row, time_row)
 
-        availabilities: Dict[str, List[TimeSlot]] = defaultdict(list)
+        availabilities: Mapping[str, List[TimeSlot]] = defaultdict(list)
         for row in reader:
             name = row[0]
 
@@ -134,44 +125,5 @@ def parse_doodle_to_stub(filename: str) -> Tuple[Iterable[str], Iterable[TimeSlo
         return tutors, sessions
 
 
-def parse_doodle_hack(filename: str, tutors: Iterable[Tutor],
-                      sessions: Iterable[Session]) \
-        -> Dict[Tuple[Tutor, Session], bool]:
-    """Parse a Doodle CSV to create a dictionary that maps a tutor name
-    to a list of all their available time slots.
-
-    Hacked version to support the way the allocation engine requires.
-    """
-    tutor_map = {tutor.name: tutor for tutor in tutors}
-    session_map = {TimeSlot(session.day.value, session.start_time,
-                            session.duration): session for session in sessions}
-
-    with open(filename, 'r') as file:
-        reader = iter(csv.reader(file))
-
-        # skip the first 4 rows since it is just doodle garbage
-        reader = islice(reader, 4, None)
-
-        day_row = next(reader)
-        time_row = next(reader)
-        days = _assign_columns_timeslots(day_row, time_row)
-
-        availabilities: Dict[Tuple[Tutor, Session], bool] = {}
-        for row in reader:
-            name = row[0]
-
-            # last row is always a count of availabilities for a timeslot
-            if name == "Count":
-                break
-
-            # add every availability timeslot
-            for column, status in islice(enumerate(row), 1, None):
-                tutor = tutor_map[name]
-                session = session_map[days[column]]
-                availabilities[(tutor, session)] = status == "OK"
-
-        return availabilities
-
-
 if __name__ == "__main__":
     print(parse_doodle("sample_doodle.csv"))
diff --git a/allocate/model.py b/allocate/model.py
index ceea073..c928e6c 100644
--- a/allocate/model.py
+++ b/allocate/model.py
@@ -2,6 +2,17 @@ from enum import Enum
 from dataclasses import dataclass
 
 
+@dataclass(eq=True, frozen=True)
+class TimeSlot:
+    """Data class used to represent a class that occupied a time slot.
+
+    This is practically used as a Session when the id of the session is unknown.
+    """
+    day: str
+    start: int
+    duration: int
+
+
 @dataclass(eq=True, frozen=True)
 class Tutor:
     name: str
diff --git a/sample_availability.csv b/sample_availability.csv
new file mode 100644
index 0000000..c52ac09
--- /dev/null
+++ b/sample_availability.csv
@@ -0,0 +1,6 @@
+,Tue,Tue,Tue,Wed,Wed,Wed,Thu,Fri,Fri
+,11,13,15,8,9,12,9,11,12
+,2,2,2,1,1,2,2,1,1
+Brae,1,,1,,1,1,1,,1
+Henry,1,1,1,1,,,1,,
+Emily,,1,,,1,1,,1,1


Test patch:
diff --git a/tests/allocation_test.py b/tests/allocation_test.py
index ec34031..8b0eb87 100644
--- a/tests/allocation_test.py
+++ b/tests/allocation_test.py
@@ -32,12 +32,27 @@ Henry,,,,,,"""
     def test_sample_allocation(self):
         """Tests to ensure standard allocation works as expected
 
-        allocate sample_tutors.csv sample_sessions.csv sample_doodle.csv
+        allocate sample_tutors.csv sample_sessions.csv sample_availability.csv
         """
         output = io.StringIO()
         with redirect_stdout(output):
             run_allocation("sample_tutors.csv", "sample_sessions.csv",
-                           "sample_doodle.csv")
+                           "sample_availability.csv")
+
+        self.assertEqual(output.getvalue().split(),
+                         AllocationTest.EXPECTED_ALLOCATION.split(),
+                         "Allocation output does not match")
+
+    def test_sample_allocation_doodle(self):
+        """Tests to ensure standard allocation works as expected with
+        a doodle file
+
+        allocate --doodle sample_tutors.csv sample_sessions.csv sample_doodle.csv
+        """
+        output = io.StringIO()
+        with redirect_stdout(output):
+            run_allocation("sample_tutors.csv", "sample_sessions.csv",
+                           "sample_doodle.csv", doodle=True)
 
         self.assertEqual(output.getvalue().split(),
                          AllocationTest.EXPECTED_ALLOCATION.split(),
@@ -47,9 +62,9 @@ Henry,,,,,,"""
         """Tests to ensure that stubs are generated from availability
         appropriately
 
-        allocate --stub stub_tutors.csv stub_sessions.csv sample_doodle.csv
+        allocate --stub stub_tutors.csv stub_sessions.csv sample_availability.csv
         """
-        stub_files("stub_tutors.csv", "stub_sessions.csv", "sample_doodle.csv")
+        stub_files("stub_tutors.csv", "stub_sessions.csv", "sample_availability.csv")
 
         with open("stub_tutors.csv") as file:
             self.assertEqual([line.strip() for line in file.readlines()],
@@ -61,7 +76,28 @@ Henry,,,,,,"""
                              AllocationTest.EXPECTED_SESSIONS_STUB.split(),
                              "Session stub file does not match")
 
+    def test_stub_generation_doodle(self):
+        """Tests to ensure that stubs are generated from availability
+        appropriately from doodle files
+
+        allocate --stub --doodle stub_tutors.csv stub_sessions.csv sample_doodle.csv
+        """
+        stub_files("stub_tutors_doodle.csv", "stub_sessions_doodle.csv",
+                   "sample_doodle.csv", doodle=True)
+
+        with open("stub_tutors_doodle.csv") as file:
+            self.assertEqual([line.strip() for line in file.readlines()],
+                             AllocationTest.EXPECTED_TUTORS_STUB.split(),
+                             "Tutor stub file does not match")
+
+        with open("stub_sessions_doodle.csv") as file:
+            self.assertEqual([line.strip() for line in file.readlines()],
+                             AllocationTest.EXPECTED_SESSIONS_STUB.split(),
+                             "Session stub file does not match")
+
     @classmethod
     def tearDownClass(cls):
         os.remove("stub_tutors.csv")
+        os.remove("stub_tutors_doodle.csv")
         os.remove("stub_sessions.csv")
+        os.remove("stub_sessions_doodle.csv")