From 1778c5fe62ff0e82598ec1fd429b1b0597743c6c Mon Sep 17 00:00:00 2001
From: Hanne Moa <hanne.moa@sikt.no>
Date: Wed, 28 Jun 2023 11:33:32 +0200
Subject: [PATCH] Fix some mypy complaints

---
 src/zinolib/compat.py      |  1 +
 src/zinolib/event_types.py | 11 ++++++-----
 src/zinolib/zino1.py       | 34 ++++++++++++++++++++++++----------
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/src/zinolib/compat.py b/src/zinolib/compat.py
index 5a245fe..b7c08d7 100644
--- a/src/zinolib/compat.py
+++ b/src/zinolib/compat.py
@@ -1,3 +1,4 @@
+# mypy: ignore-errors
 try:
     from enum import StrEnum
 except ImportError:
diff --git a/src/zinolib/event_types.py b/src/zinolib/event_types.py
index 8da6138..0d03b6e 100644
--- a/src/zinolib/event_types.py
+++ b/src/zinolib/event_types.py
@@ -1,5 +1,5 @@
 from datetime import datetime, timedelta, timezone
-from typing import Optional, ClassVar, List, TypeVar, Union, Dict
+from typing import Optional, ClassVar, List, TypeVar, Union, Dict, Generic
 
 from pydantic import IPvAnyAddress
 from pydantic import BaseModel
@@ -281,19 +281,20 @@ class EventEngine:
     A list of already existing events can be manipulated by an instance of
     this class out of the box, but the actual IO is done by subclasses.
     """
-    events: Dict[str, EventType]
+    events: Dict[int, Event]
 
     def __init__(self, session=None):
         self.session = session
         self.events = {}
 
-    def _get_event(self, event_or_id: EventOrId) -> EventType:
-        if isinstance(event_or_id, EventType):
+    def _get_event(self, event_or_id: EventOrId) -> Event:
+        if isinstance(event_or_id, Event):
             return event_or_id
         if isinstance(event_or_id, int):
             return self.events[event_or_id]
+        raise ValueError("Unknown type")
 
-    def _set_event(self, event: EventType):
+    def _set_event(self, event: Event):
         self.events[event.id] = event
 
     def check_session(self):
diff --git a/src/zinolib/zino1.py b/src/zinolib/zino1.py
index 976c7b3..aa417d4 100644
--- a/src/zinolib/zino1.py
+++ b/src/zinolib/zino1.py
@@ -48,11 +48,23 @@
 """
 
 from datetime import datetime, timezone
-from typing import List
+from typing import List, Dict, Union, TypedDict
 
 from .event_types import EventType, Event, EventEngine, HistoryEntry, LogEntry
 
 
+HistoryDict = TypedDict(
+    "HistoryDict",
+    {"date": datetime, "user": str, "log": str},
+    total=False,
+)
+LogDict = TypedDict(
+    "LogDict",
+    {"log": str, "date": datetime},
+    total=False,
+)
+
+
 def convert_timestamp(timestamp: int) -> datetime:
     return datetime.fromtimestamp(timestamp, timezone.utc)
 
@@ -134,7 +146,7 @@ def parse_response(cls, history_data: List[str]):
             ..
         ]
         """
-        history_list = []
+        history_list: List[HistoryDict] = []
         for row in history_data:
             if row == " ":  # end of history
                 continue
@@ -142,12 +154,14 @@ def parse_response(cls, history_data: List[str]):
                 history_list[-1]["log"] += row.strip()
             else:
                 timestamp, body = row.split(" ", 1)
-                dt = convert_timestamp(timestamp)
-                entry = {"date": dt}
+                dt = convert_timestamp(int(timestamp))
+                entry: HistoryDict = {"date": dt}
                 if " " in body:  # server generated
-                    entry = {"log": body, "user": cls.SYSTEM_USER}
+                    entry["user"] = cls.SYSTEM_USER
+                    entry["log"] = body
                 else:
-                    entry = {"log": "", "user": body}
+                    entry["user"] =  body
+                    entry["log"] = ""
                 history_list.append(entry)
 
         return history_list
@@ -169,7 +183,7 @@ def get_log(session, event_id: int):
         return session.get_raw_log(event_id).data
 
     @staticmethod
-    def parse_response(log_data: List[dict]):
+    def parse_response(log_data: List[str]) -> List[LogDict]:
         """
         Input:
         [
@@ -186,10 +200,10 @@ def parse_response(log_data: List[dict]):
             ..
         ]
         """
-        log_list = []
+        log_list: List[LogDict] = []
         for row in log_data:
             timestamp, log = row.split(" ", 1)
-            dt = convert_timestamp(timestamp)
+            dt = convert_timestamp(int(timestamp))
             log_list.append({"date": dt, "log": log})
         return log_list
 
@@ -232,5 +246,5 @@ def get_history_for_id(self, event_id: int):
     def get_log_for_id(self, event_id: int):
         self.check_session()
         raw_log = self._log_adapter.get_log(self.session, event_id)
-        parsed_log = self._log_adapter.parse_log_response(raw_log)
+        parsed_log = self._log_adapter.parse_response(raw_log)
         return LogEntry.create_list(parsed_log)