Fix linting in tests

Author: gongcastro

Makefile

@@ -2,4 +2,4 @@ test:
 	uvx ruff format
 	uvx ruff check --fix .
 	uvx ty check
-	uvx pytest -v --exitfirst
+	uv run pytest -v --exitfirst

babylab/utils.py

@@ -194,14 +194,14 @@ def count_col(
     return counts
 
 
-def get_year_weeks(year: int) -> Generator[datetime]:
+def get_year_weeks(year: int) -> Generator[datetime, datetime, datetime]:
     """Get week numbers of the year.
 
     Args:
         year (int): Year to get weeks for.
 
     Yields:
-        int: Number of weeks in the year.
+        Generator[datetime, datetime, datetime]: Number of weeks in the year.
     """
     date_first = date(year, 1, 1)
     date_first += timedelta(days=6 - date_first.weekday())

tests/test_utils.py

@@ -119,7 +119,11 @@ def test_get_apt_table_id_list(apt_id: str | Sequence[str] | None = None, k: int
         apt_id = [apt_id]
 
     ppt_id = set(i.split(":")[0] for i in apt_id)
-    df = utils.get_apt_table(conftest.RECORDS, ppt_id=sample(ppt_id, k=k))
+
+    if k > len(ppt_id):
+        k = len(ppt_id)
+
+    df = utils.get_apt_table(conftest.RECORDS, ppt_id=sample(list(ppt_id), k=k))
 
     assert isinstance(df, pl.DataFrame)
     assert all(p in ppt_id for p in df["record_id"].unique().to_list())
@@ -148,7 +152,11 @@ def test_get_que_table_id_list(que_id: str | list[str] | None = None, k: int = 1
         que_id = [que_id]
 
     ppt_id = set(i.split(":")[0] for i in que_id)
-    df = utils.get_apt_table(conftest.RECORDS, ppt_id=sample(ppt_id, k=k))
+
+    if k > len(ppt_id):
+        k = len(ppt_id)
+
+    df = utils.get_apt_table(conftest.RECORDS, ppt_id=sample(list(ppt_id), k=k))
     assert isinstance(df, pl.DataFrame)
     assert all(p in ppt_id for p in df["record_id"].unique().to_list())