use native types List->list, Dict->dict, Set->set, Tuple->tuple

reasoning_gym/graphs/course_schedule.py

@@ -8,7 +8,7 @@ https://leetcode.com/problems/course-schedule/description/
 from collections import defaultdict
 from dataclasses import dataclass
 from random import Random
-from typing import List, Optional
+from typing import Optional
 
 from ..factory import ProceduralDataset, register_dataset
 
@@ -54,7 +54,7 @@ class CourseScheduleDataset(ProceduralDataset):
     def __init__(self, config: CourseScheduleConfig):
         super().__init__(config=config, seed=config.seed, size=config.size)
 
-    def _can_finish(self, num_courses: int, prerequisites: List[List[int]]) -> bool:
+    def _can_finish(self, num_courses: int, prerequisites: list[list[int]]) -> bool:
         adj = defaultdict(list)
         for course, prereq in prerequisites:
             adj[course].append(prereq)
@@ -81,7 +81,7 @@ class CourseScheduleDataset(ProceduralDataset):
 
         return True
 
-    def _create_prerequisites(self, rng: Random, courses: List[int], solvable: bool) -> List[List[int]]:
+    def _create_prerequisites(self, rng: Random, courses: list[int], solvable: bool) -> list[list[int]]:
         """Create a list of prerequisites for each course"""
         prerequisites = []
         # Generate a valid course schedule

reasoning_gym/graphs/family_relationships.py

@@ -2,7 +2,7 @@ import random
 from dataclasses import dataclass, field
 from enum import StrEnum
 from itertools import count
-from typing import List, Optional, Set, Tuple
+from typing import Optional
 
 from ..factory import ProceduralDataset, register_dataset
 
@@ -37,8 +37,8 @@ class Person:
     gender: Gender
     id: int
     spouse: Optional["Person"] = None
-    parents: List["Person"] = field(default_factory=list)
-    children: List["Person"] = field(default_factory=list)
+    parents: list["Person"] = field(default_factory=list)
+    children: list["Person"] = field(default_factory=list)
 
     def __hash__(self):
         return self.id
@@ -65,7 +65,7 @@ class FamilyRelationshipsConfig:
 
     min_family_size: int = 4
     max_family_size: int = 8
-    male_names: List[str] = field(
+    male_names: list[str] = field(
         default_factory=lambda: [
             "James",
             "John",
@@ -112,7 +112,7 @@ class FamilyRelationshipsConfig:
             "Finn",
         ]
     )
-    female_names: List[str] = field(
+    female_names: list[str] = field(
         default_factory=lambda: [
             "Mary",
             "Patricia",
@@ -207,7 +207,7 @@ class FamilyRelationshipsDataset(ProceduralDataset):
             },
         }
 
-    def _generate_family(self, rng: random.Random) -> Set[Person]:
+    def _generate_family(self, rng: random.Random) -> set[Person]:
         """Generate a random family tree"""
         family_size = rng.randint(self.config.min_family_size, self.config.max_family_size)
         family = set()
@@ -292,8 +292,8 @@ class FamilyRelationshipsDataset(ProceduralDataset):
         return family
 
     def _get_relationship_question(
-        self, rng: random.Random, family: Set[Person]
-    ) -> Tuple[Person, Person, Relationship]:
+        self, rng: random.Random, family: set[Person]
+    ) -> tuple[Person, Person, Relationship]:
         """Select two family members and determine their relationship"""
         person1, person2 = rng.sample(list(family), 2)
 
@@ -326,7 +326,7 @@ class FamilyRelationshipsDataset(ProceduralDataset):
 
         return person1, person2, relationship
 
-    def _generate_story(self, family: Set[Person]) -> str:
+    def _generate_story(self, family: set[Person]) -> str:
         """Generate a story describing the family relationships"""
         story_parts = []
 

reasoning_gym/graphs/largest_island.py

@@ -7,7 +7,7 @@ https://leetcode.com/problems/max-area-of-island/description/
 from collections import deque
 from dataclasses import dataclass
 from random import Random
-from typing import List, Optional
+from typing import Optional
 
 from ..factory import ProceduralDataset, register_dataset
 
@@ -58,7 +58,7 @@ class LargestIslandDataset(ProceduralDataset):
     def _is_valid_cell(self, r: int, c: int) -> bool:
         return 0 <= r < self.config.rows and 0 <= c < self.config.cols
 
-    def _create_grid(self, rng: Random) -> List[List[int]]:
+    def _create_grid(self, rng: Random) -> list[list[int]]:
         """Create a random grid of islands using a random walk algorithm"""
         grid = [[0] * self.config.cols for _ in range(self.config.rows)]
         directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # Up, Down, Left, Right
@@ -81,7 +81,7 @@ class LargestIslandDataset(ProceduralDataset):
 
         return grid
 
-    def _get_largest_island(self, grid: List[List[int]]) -> int:
+    def _get_largest_island(self, grid: list[list[int]]) -> int:
         """Find the largest island in the grid"""
         directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # Up, Down, Left, Right
         visited = set()
@@ -108,11 +108,11 @@ class LargestIslandDataset(ProceduralDataset):
 
         return max_area
 
-    def _grid_to_string(self, grid: List[List[int]]) -> str:
+    def _grid_to_string(self, grid: list[list[int]]) -> str:
         """Convert grid to a string representation"""
         return "\n".join(" ".join(str(cell) for cell in row) for row in grid)
 
-    def _string_to_board(self, grid_str: str) -> List[List[int]]:
+    def _string_to_board(self, grid_str: str) -> list[list[int]]:
         """Convert string representation to a grid"""
         return [[int(cell) for cell in row.split()] for row in grid_str.split("\n")]
 

reasoning_gym/graphs/quantum_lock.py

@@ -164,7 +164,7 @@ Buttons:
         # If no solution found, regenerate
         return self.generate_quantum_puzzle(rng, difficulty)
 
-    def score_answer(self, answer: Optional[str], entry: dict[str, any]) -> float:
+    def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
         """Determine if the solution provided solves the task.
 
         The function awards 1.0 for a correct answer and less otherwise.