From f1ada62c800ce455c752b0137da37e7e5c9eb44c Mon Sep 17 00:00:00 2001
From: Allen Porter <allen.porter@gmail.com>
Date: Sat, 3 Aug 2024 22:37:28 +0000
Subject: [PATCH] Improve typing

---
 home_assistant_datasets/tool/assist/conftest.py | 2 +-
 script/export.py                                | 2 +-
 script/human_eval_metrics.py                    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/home_assistant_datasets/tool/assist/conftest.py b/home_assistant_datasets/tool/assist/conftest.py
index 0c318f29..99d81445 100644
--- a/home_assistant_datasets/tool/assist/conftest.py
+++ b/home_assistant_datasets/tool/assist/conftest.py
@@ -74,7 +74,7 @@ def pytest_generate_tests(metafunc: Any) -> None:
 
 
 @pytest.fixture(autouse=True)
-def restore_tz() -> Generator[None, None]:
+def restore_tz() -> Generator[None, None, None]:
     yield
     # Home Assistant teardown seems to run too soon and expects this so try to
     # patch it in first.
diff --git a/script/export.py b/script/export.py
index 07b1e1f3..c865708f 100644
--- a/script/export.py
+++ b/script/export.py
@@ -5,7 +5,7 @@
 import sys
 
 
-def main():
+def main() -> None:
     lines = sys.stdin.readlines()
     for line in lines:
         record = json.loads(line)
diff --git a/script/human_eval_metrics.py b/script/human_eval_metrics.py
index 21581b10..1c27098f 100644
--- a/script/human_eval_metrics.py
+++ b/script/human_eval_metrics.py
@@ -94,7 +94,7 @@ def main() -> None:
         print(yaml.dump(model_results, sort_keys=True, explicit_start=True))
 
         if args.samples:
-            limited_samples = {}
+            limited_samples: dict[str, dict[str, str]] = {}
             for model_id in model_samples:
                 limited_samples[model_id] = {}
                 for label in model_samples[model_id]: