Refactor/typings auto discovery (#831)

* refactor: substituting Option for union None type

* refactor: replaced Dict with dict

* refactor: added missing generic types and type annotation for dicts that infer to unknown

---------

Co-authored-by: burtenshaw <ben.burtenshaw@gmail.com>

Author: hinogi

src/openenv/auto/_discovery.py

@@ -17,6 +17,8 @@
 This enables AutoEnv to work without coupling to src/envs/ directory.
 """
 
+from __future__ import annotations
+
 import importlib
 import importlib.metadata
 import importlib.resources
@@ -26,7 +28,7 @@
 import tempfile
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, Dict, Optional, Type
+from typing import Any, Type
 
 import yaml
 
@@ -63,10 +65,10 @@ class EnvironmentInfo:
     action_class_name: str
     observation_class_name: str
     default_image: str
-    spec_version: Optional[int] = None
-    manifest: Optional[Dict[str, Any]] = None
+    spec_version: int | None = None
+    manifest: dict[str, Any] | None = None
 
-    def get_client_class(self) -> Type:
+    def get_client_class(self) -> Type[Any]:
         """
         Dynamically import and return the client class.
 
@@ -90,7 +92,7 @@ def get_client_class(self) -> Type:
                 f"Class {self.client_class_name} not found in {self.client_module_path}: {e}"
             ) from e
 
-    def get_action_class(self) -> Type:
+    def get_action_class(self) -> Type[Any]:
         """
         Dynamically import and return the action class.
 
@@ -114,7 +116,7 @@ def get_action_class(self) -> Type:
                 f"Class {self.action_class_name} not found in {self.client_module_path}: {e}"
             ) from e
 
-    def get_observation_class(self) -> Type:
+    def get_observation_class(self) -> Type[Any]:
         """
         Dynamically import and return the observation class.
 
@@ -225,7 +227,7 @@ def _infer_class_name(env_name: str, class_type: str) -> str:
 
 def _load_manifest_from_package(
     package_name: str, module_name: str
-) -> Optional[Dict[str, Any]]:
+) -> dict[str, Any] | None:
     """
     Load openenv.yaml manifest from an installed package.
 
@@ -259,7 +261,7 @@ def _load_manifest_from_package(
 
 def _create_env_info_from_package(
     package_name: str, module_name: str, version: str
-) -> Optional[EnvironmentInfo]:
+) -> EnvironmentInfo | None:
     """
     Create EnvironmentInfo from an installed package.
 
@@ -347,17 +349,17 @@ class EnvironmentDiscovery:
 
     def __init__(self):
         """Initialize discovery system."""
-        self._cache: Optional[Dict[str, EnvironmentInfo]] = None
+        self._cache: dict[str, EnvironmentInfo] | None = None
         self._cache_file = Path(tempfile.gettempdir()) / "openenv_discovery_cache.json"
 
-    def _discover_installed_packages(self) -> Dict[str, EnvironmentInfo]:
+    def _discover_installed_packages(self) -> dict[str, EnvironmentInfo]:
         """
         Discover all installed openenv-* packages.
 
         Returns:
             Dictionary mapping env_key to EnvironmentInfo
         """
-        environments = {}
+        environments: dict[str, EnvironmentInfo] = {}
 
         # Invalidate import caches to ensure we pick up newly installed packages
         importlib.invalidate_caches()
@@ -403,7 +405,7 @@ def _discover_installed_packages(self) -> Dict[str, EnvironmentInfo]:
 
         return environments
 
-    def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]:
+    def _load_cache(self) -> dict[str, EnvironmentInfo] | None:
         """
         Load cached discovery results.
 
@@ -418,7 +420,7 @@ def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]:
                 cache_data = json.load(f)
 
             # Reconstruct EnvironmentInfo objects
-            cache = {}
+            cache: dict[str, EnvironmentInfo] = {}
             for env_key, env_data in cache_data.items():
                 cache[env_key] = EnvironmentInfo(**env_data)
 
@@ -427,7 +429,7 @@ def _load_cache(self) -> Optional[Dict[str, EnvironmentInfo]]:
             logger.warning(f"Failed to load discovery cache: {e}")
             return None
 
-    def _save_cache(self, environments: Dict[str, EnvironmentInfo]) -> None:
+    def _save_cache(self, environments: dict[str, EnvironmentInfo]) -> None:
         """
         Save discovery results to cache.
 
@@ -445,7 +447,7 @@ def _save_cache(self, environments: Dict[str, EnvironmentInfo]) -> None:
         except Exception as e:
             logger.warning(f"Failed to save discovery cache: {e}")
 
-    def discover(self, use_cache: bool = True) -> Dict[str, EnvironmentInfo]:
+    def discover(self, use_cache: bool = True) -> dict[str, EnvironmentInfo]:
         """
         Discover all installed OpenEnv environments.
 
@@ -481,7 +483,7 @@ def discover(self, use_cache: bool = True) -> Dict[str, EnvironmentInfo]:
 
         return environments
 
-    def get_environment(self, env_key: str) -> Optional[EnvironmentInfo]:
+    def get_environment(self, env_key: str) -> EnvironmentInfo | None:
         """
         Get information about a specific environment.
 
@@ -500,7 +502,7 @@ def get_environment(self, env_key: str) -> Optional[EnvironmentInfo]:
         environments = self.discover()
         return environments.get(env_key)
 
-    def get_environment_by_name(self, name: str) -> Optional[EnvironmentInfo]:
+    def get_environment_by_name(self, name: str) -> EnvironmentInfo | None:
         """
         Get environment info by flexible name matching.
 
@@ -554,7 +556,7 @@ def clear_cache(self) -> None:
 
 
 # Global discovery instance
-_global_discovery: Optional[EnvironmentDiscovery] = None
+_global_discovery: EnvironmentDiscovery | None = None
 
 
 def get_discovery() -> EnvironmentDiscovery: